Update app.py

app.py

@@ -127,20 +127,9 @@ ELEVATION_MAP = {
 }
 
 def snap_to_nearest(value, options):
-    """Snap a value to the nearest option in a list."""
     return min(options, key=lambda x: abs(x - value))
 
 def build_lighting_prompt(azimuth: float, elevation: float) -> str:
-    """
-    Build a lighting prompt from azimuth and elevation values.
-    
-    Args:
-        azimuth: Horizontal rotation in degrees (0-360)
-        elevation: Vertical angle in degrees (-90 to 90)
-    
-    Returns:
-        Formatted prompt string for the LoRA
-    """
     azimuth_snapped = snap_to_nearest(azimuth, list(AZIMUTH_MAP.keys()))
     elevation_snapped = snap_to_nearest(elevation, list(ELEVATION_MAP.keys()))
     
@@ -161,9 +150,6 @@ def infer_lighting_edit(
     height: int = 1024,
     width: int = 1024,
 ):
-    """
-    Edit the lighting of an image using Qwen Image Edit 2511 with multi-angle lighting LoRA.
-    """
     global loaded
     progress = gr.Progress(track_tqdm=True)
     
@@ -199,9 +185,7 @@ def infer_lighting_edit(
 def update_dimensions_on_upload(image):
     if image is None:
         return 1024, 1024
-    
     original_width, original_height = image.size
-    
     if original_width > original_height:
         new_width = 1024
         aspect_ratio = original_height / original_width
@@ -210,18 +194,11 @@ def update_dimensions_on_upload(image):
         new_height = 1024
         aspect_ratio = original_width / original_height
         new_width = int(new_height * aspect_ratio)
-        
     new_width = (new_width // 8) * 8
     new_height = (new_height // 8) * 8
-    
     return new_width, new_height
 
 class LightingControl3D(gr.HTML):
-    """
-    A 3D lighting control component using Three.js.
-    Outputs: { azimuth: number, elevation: number }
-    Accepts imageUrl prop to display user's uploaded image on the plane.
-    """
     def __init__(self, value=None, imageUrl=None, **kwargs):
         if value is None:
             value = {"azimuth": 0, "elevation": 0}
@@ -237,14 +214,12 @@ class LightingControl3D(gr.HTML):
             const wrapper = element.querySelector('#lighting-control-wrapper');
             const promptOverlay = element.querySelector('#prompt-overlay');
             
-            // Wait for THREE to load
             const initScene = () => {
                 if (typeof THREE === 'undefined') {
                     setTimeout(initScene, 100);
                     return;
                 }
                 
-                // Scene setup
                 const scene = new THREE.Scene();
                 scene.background = new THREE.Color(0x1a1a1a);
                 
@@ -259,10 +234,8 @@ class LightingControl3D(gr.HTML):
                 renderer.shadowMap.type = THREE.PCFSoftShadowMap;
                 wrapper.insertBefore(renderer.domElement, promptOverlay);
                 
-                // Lighting for the scene
                 scene.add(new THREE.AmbientLight(0xffffff, 0.1));
                 
-                // Ground plane for shadows
                 const ground = new THREE.Mesh(
                     new THREE.PlaneGeometry(10, 10),
                     new THREE.ShadowMaterial({ opacity: 0.3 })
@@ -272,20 +245,16 @@ class LightingControl3D(gr.HTML):
                 ground.receiveShadow = true;
                 scene.add(ground);
                 
-                // Grid
                 scene.add(new THREE.GridHelper(8, 16, 0x333333, 0x222222));
                 
-                // Constants
                 const CENTER = new THREE.Vector3(0, 0.75, 0);
                 const BASE_DISTANCE = 2.5;
                 const AZIMUTH_RADIUS = 2.4;
                 const ELEVATION_RADIUS = 1.8;
                 
-                // State
                 let azimuthAngle = props.value?.azimuth || 0;
                 let elevationAngle = props.value?.elevation || 0;
                 
-                // Mappings
                 const azimuthSteps = [0, 45, 90, 135, 180, 225, 270, 315];
                 const elevationSteps = [-90, 0, 90];
                 const azimuthNames = {
@@ -299,7 +268,6 @@ class LightingControl3D(gr.HTML):
                     return steps.reduce((prev, curr) => Math.abs(curr - value) < Math.abs(prev - value) ? curr : prev);
                 }
                 
-                // Create placeholder texture (smiley face)
                 function createPlaceholderTexture() {
                     const canvas = document.createElement('canvas');
                     canvas.width = 256;
@@ -324,7 +292,6 @@ class LightingControl3D(gr.HTML):
                     return new THREE.CanvasTexture(canvas);
                 }
                 
-                // Target image plane
                 let currentTexture = createPlaceholderTexture();
                 const planeMaterial = new THREE.MeshStandardMaterial({ map: currentTexture, side: THREE.DoubleSide, roughness: 0.5, metalness: 0 });
                 let targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.2, 1.2), planeMaterial);
@@ -332,13 +299,10 @@ class LightingControl3D(gr.HTML):
                 targetPlane.receiveShadow = true;
                 scene.add(targetPlane);
                 
-                // Function to update texture from image URL
                 function updateTextureFromUrl(url) {
                     if (!url) {
-                        // Reset to placeholder
                         planeMaterial.map = createPlaceholderTexture();
                         planeMaterial.needsUpdate = true;
-                        // Reset plane to square
                         scene.remove(targetPlane);
                         targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.2, 1.2), planeMaterial);
                         targetPlane.position.copy(CENTER);
@@ -355,7 +319,6 @@ class LightingControl3D(gr.HTML):
                         planeMaterial.map = texture;
                         planeMaterial.needsUpdate = true;
                         
-                        // Adjust plane aspect ratio to match image
                         const img = texture.image;
                         if (img && img.width && img.height) {
                             const aspect = img.width / img.height;
@@ -382,63 +345,51 @@ class LightingControl3D(gr.HTML):
                     });
                 }
                 
-                // Check for initial imageUrl
                 if (props.imageUrl) {
                     updateTextureFromUrl(props.imageUrl);
                 }
                 
-                // --- NEW LIGHT MODEL: STUDIO SQUARE WITH RAYS ---
                 const lightGroup = new THREE.Group();
 
-                // 1. Studio Housing (Black/Dark Grey)
-                const housingGeo = new THREE.BoxGeometry(0.7, 0.7, 0.2);
+                const housingGeo = new THREE.BoxGeometry(0.8, 0.8, 0.2);
                 const housingMat = new THREE.MeshStandardMaterial({ 
-                    color: 0x222222, 
-                    roughness: 0.3,
-                    metalness: 0.7
+                    color: 0x111111,   
+                    roughness: 0.7,
+                    metalness: 0.5
                 });
                 const housing = new THREE.Mesh(housingGeo, housingMat);
                 
-                // 2. Light Surface (Bright White Square)
-                const faceGeo = new THREE.PlaneGeometry(0.65, 0.65);
-                const faceMat = new THREE.MeshStandardMaterial({ 
-                    color: 0xffffff, 
+                const panelGeo = new THREE.PlaneGeometry(0.75, 0.75);
+                const panelMat = new THREE.MeshStandardMaterial({ 
+                    color: 0xffffff,
                     emissive: 0xffffff,
-                    emissiveIntensity: 3.0, 
-                    roughness: 0.0
+                    emissiveIntensity: 2.0, 
+                    roughness: 0.1
                 });
-                const face = new THREE.Mesh(faceGeo, faceMat);
-                face.position.z = 0.105; // Slightly in front of housing
-                housing.add(face);
-                
-                // 3. Light Rays (Volumetric Beam Effect)
-                // Using a transparent cylinder to simulate a beam
-                const beamHeight = 4.0;
-                // Cylinder top (0.5) is light size, bottom (1.2) is spread
-                const beamGeo = new THREE.CylinderGeometry(0.5, 1.2, beamHeight, 32, 1, true);
+                const panel = new THREE.Mesh(panelGeo, panelMat);
+                panel.position.z = 0.101; 
                 
-                // Shift geometry so narrow end starts at 0 and extends positive
-                beamGeo.translate(0, -beamHeight / 2, 0);
-                beamGeo.rotateX(-Math.PI / 2); // Rotate to point along +Z axis
+                housing.add(panel);
                 
+                const beamGeo = new THREE.ConeGeometry(1.2, 5, 32, 1, true);
                 const beamMat = new THREE.MeshBasicMaterial({
-                    color: 0xffffff,
+                    color: 0xffffee,
                     transparent: true,
-                    opacity: 0.06,
+                    opacity: 0.08,
                     side: THREE.DoubleSide,
                     depthWrite: false,
                     blending: THREE.AdditiveBlending
                 });
+                const lightBeam = new THREE.Mesh(beamGeo, beamMat);
                 
-                const beam = new THREE.Mesh(beamGeo, beamMat);
-                beam.position.z = 0.12; // Start just in front of the face
-                housing.add(beam);
-                
+                lightBeam.rotation.x = Math.PI / 2; 
+                lightBeam.position.z = 2.5; 
+                housing.add(lightBeam);
+
                 lightGroup.add(housing);
                 
-                // Actual Light Source
-                const spotLight = new THREE.SpotLight(0xffffff, 15, 20, Math.PI / 3, 0.5, 1);
-                spotLight.position.set(0, 0, 0.1);
+                const spotLight = new THREE.SpotLight(0xffffff, 10, 10, Math.PI / 3, 1, 1);
+                spotLight.position.set(0, 0, 0.1); 
                 spotLight.castShadow = true;
                 spotLight.shadow.mapSize.width = 1024;
                 spotLight.shadow.mapSize.height = 1024;
@@ -454,9 +405,6 @@ class LightingControl3D(gr.HTML):
                 
                 scene.add(lightGroup);
                 
-                // --- CONTROLS: COLORS UPDATED ---
-                
-                // Azimuth ring (YELLOW)
                 const azimuthRing = new THREE.Mesh(
                     new THREE.TorusGeometry(AZIMUTH_RADIUS, 0.04, 16, 64),
                     new THREE.MeshStandardMaterial({ color: 0xffff00, emissive: 0xffff00, emissiveIntensity: 0.3 })
@@ -465,7 +413,6 @@ class LightingControl3D(gr.HTML):
                 azimuthRing.position.y = 0.05;
                 scene.add(azimuthRing);
                 
-                // Azimuth Handle (YELLOW)
                 const azimuthHandle = new THREE.Mesh(
                     new THREE.SphereGeometry(0.18, 16, 16),
                     new THREE.MeshStandardMaterial({ color: 0xffff00, emissive: 0xffff00, emissiveIntensity: 0.5 })
@@ -473,7 +420,6 @@ class LightingControl3D(gr.HTML):
                 azimuthHandle.userData.type = 'azimuth';
                 scene.add(azimuthHandle);
                 
-                // Elevation arc (BLUE)
                 const arcPoints = [];
                 for (let i = 0; i <= 32; i++) {
                     const angle = THREE.MathUtils.degToRad(-90 + (180 * i / 32));
@@ -486,7 +432,6 @@ class LightingControl3D(gr.HTML):
                 );
                 scene.add(elevationArc);
                 
-                // Elevation Handle (BLUE)
                 const elevationHandle = new THREE.Mesh(
                     new THREE.SphereGeometry(0.18, 16, 16),
                     new THREE.MeshStandardMaterial({ color: 0x0000ff, emissive: 0x0000ff, emissiveIntensity: 0.5 })
@@ -494,13 +439,12 @@ class LightingControl3D(gr.HTML):
                 elevationHandle.userData.type = 'elevation';
                 scene.add(elevationHandle);
                 
-                // --- REFRESH BUTTON ---
                 const refreshBtn = document.createElement('button');
                 refreshBtn.innerHTML = 'Reset View';
                 refreshBtn.style.position = 'absolute';
                 refreshBtn.style.top = '15px';
                 refreshBtn.style.right = '15px';
-                refreshBtn.style.background = '#e63e00'; // Theme primary color
+                refreshBtn.style.background = '#e63e00';
                 refreshBtn.style.color = '#fff';
                 refreshBtn.style.border = 'none';
                 refreshBtn.style.padding = '8px 16px';
@@ -540,7 +484,6 @@ class LightingControl3D(gr.HTML):
                     azimuthHandle.position.set(AZIMUTH_RADIUS * Math.sin(azRad), 0.05, AZIMUTH_RADIUS * Math.cos(azRad));
                     elevationHandle.position.set(-0.8, ELEVATION_RADIUS * Math.sin(elRad) + CENTER.y, ELEVATION_RADIUS * Math.cos(elRad));
                     
-                    // Update prompt
                     const azSnap = snapToNearest(azimuthAngle, azimuthSteps);
                     const elSnap = snapToNearest(elevationAngle, elevationSteps);
                     let prompt = 'Light source from';
@@ -560,7 +503,6 @@ class LightingControl3D(gr.HTML):
                     trigger('change', props.value);
                 }
                 
-                // Raycasting
                 const raycaster = new THREE.Raycaster();
                 const mouse = new THREE.Vector2();
                 let isDragging = false;
@@ -633,7 +575,6 @@ class LightingControl3D(gr.HTML):
                         dragTarget.material.emissiveIntensity = 0.5;
                         dragTarget.scale.setScalar(1);
                         
-                        // Snap and animate
                         const targetAz = snapToNearest(azimuthAngle, azimuthSteps);
                         const targetEl = snapToNearest(elevationAngle, elevationSteps);
                         
@@ -666,7 +607,6 @@ class LightingControl3D(gr.HTML):
                 
                 canvas.addEventListener('mouseup', onMouseUp);
                 canvas.addEventListener('mouseleave', onMouseUp);
-                // Touch support for mobile
                 canvas.addEventListener('touchstart', (e) => {
                     e.preventDefault();
                     const touch = e.touches[0];
@@ -724,24 +664,20 @@ class LightingControl3D(gr.HTML):
                     onMouseUp();
                 }, { passive: false });
                 
-                // Initial update
                 updatePositions();
                 
-                // Render loop
                 function render() {
                     requestAnimationFrame(render);
                     renderer.render(scene, camera);
                 }
                 render();
                 
-                // Handle resize
                 new ResizeObserver(() => {
                     camera.aspect = wrapper.clientWidth / wrapper.clientHeight;
                     camera.updateProjectionMatrix();
                     renderer.setSize(wrapper.clientWidth, wrapper.clientHeight);
                 }).observe(wrapper);
                 
-                // Store update functions for external calls
                 wrapper._updateFromProps = (newVal) => {
                     if (newVal && typeof newVal === 'object') {
                         azimuthAngle = newVal.azimuth ?? azimuthAngle;
@@ -752,16 +688,13 @@ class LightingControl3D(gr.HTML):
                 
                 wrapper._updateTexture = updateTextureFromUrl;
                 
-                // Watch for prop changes (imageUrl and value)
                 let lastImageUrl = props.imageUrl;
                 let lastValue = JSON.stringify(props.value);
                 setInterval(() => {
-                    // Check imageUrl changes
                     if (props.imageUrl !== lastImageUrl) {
                         lastImageUrl = props.imageUrl;
                         updateTextureFromUrl(props.imageUrl);
                     }
-                    // Check value changes (from sliders)
                     const currentValue = JSON.stringify(props.value);
                     if (currentValue !== lastValue) {
                         lastValue = currentValue;
@@ -802,7 +735,7 @@ with gr.Blocks(css=css) as demo:
             image = gr.Image(label="Input Image", type="pil", height=300)
             
             gr.Markdown("### 3D Lighting Control")
-            #gr.Markdown("*Drag the colored handles: 🟡 Azimuth (Direction), 🔵 Elevation (Height)*")
+            gr.Markdown("*Drag the colored handles: 🟡 Azimuth (Direction), 🔵 Elevation (Height)*")
             
             lighting_3d = LightingControl3D(
                 value={"azimuth": 0, "elevation": 0},
@@ -841,7 +774,7 @@ with gr.Blocks(css=css) as demo:
         with gr.Column(scale=1):
             result = gr.Image(label="Output Image", height=500)
             
-            with gr.Accordion("Advanced Settings", open=True):
+            with gr.Accordion("Advanced Settings", open=False):
                 seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                 randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                 guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
@@ -850,12 +783,10 @@ with gr.Blocks(css=css) as demo:
                 width = gr.Slider(label="Width", minimum=256, maximum=2048, step=8, value=1024)
     
     def update_prompt_from_sliders(azimuth, elevation):
-        """Update prompt preview when sliders change."""
         prompt = build_lighting_prompt(azimuth, elevation)
         return prompt
     
     def sync_3d_to_sliders(lighting_value):
-        """Sync 3D control changes to sliders."""
         if lighting_value and isinstance(lighting_value, dict):
             az = lighting_value.get('azimuth', 0)
             el = lighting_value.get('elevation', 0)
@@ -864,11 +795,9 @@ with gr.Blocks(css=css) as demo:
         return gr.update(), gr.update(), gr.update()
     
     def sync_sliders_to_3d(azimuth, elevation):
-        """Sync slider changes to 3D control."""
         return {"azimuth": azimuth, "elevation": elevation}
     
     def update_3d_image(image):
-        """Update the 3D component with the uploaded image."""
         if image is None:
             return gr.update(imageUrl=None)