Update app.py

app.py

@@ -13,7 +13,6 @@ from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
-# --- Theme Configuration ---
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -82,11 +81,10 @@ class OrangeRedTheme(Soft):
 
 orange_red_theme = OrangeRedTheme()
 
-# --- Model & Hardware Setup ---
 MAX_SEED = np.iinfo(np.int32).max
+
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2511",
     transformer=QwenImageTransformer2DModel.from_pretrained(
@@ -96,7 +94,6 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
     ),
     torch_dtype=dtype
 ).to(device)
-
 try:
     pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
     print("Flash Attention 3 Processor set successfully.")
@@ -112,18 +109,39 @@ ADAPTER_SPECS = {
 }
 loaded = False
 
-# --- Mappings ---
 AZIMUTH_MAP = {
-    0: "Front", 45: "Right Front", 90: "Right", 135: "Right Rear",
-    180: "Rear", 225: "Left Rear", 270: "Left", 315: "Left Front"
+    0: "Front",
+    45: "Right Front",
+    90: "Right",
+    135: "Right Rear",
+    180: "Rear",
+    225: "Left Rear",
+    270: "Left",
+    315: "Left Front"
 }
 
-ELEVATION_MAP = { -90: "Below", 0: "", 90: "Above" }
+ELEVATION_MAP = {
+    -90: "Below",
+    0: "",
+    90: "Above"
+}
 
 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
+    """
+    # Snap to nearest valid values
     azimuth_snapped = snap_to_nearest(azimuth, list(AZIMUTH_MAP.keys()))
     elevation_snapped = snap_to_nearest(elevation, list(ELEVATION_MAP.keys()))
     
@@ -132,7 +150,6 @@ def build_lighting_prompt(azimuth: float, elevation: float) -> str:
     else:
         return f"Light source from {ELEVATION_MAP[elevation_snapped]}"
 
-# --- Inference Function ---
 @spaces.GPU
 def infer_lighting_edit(
     image: Image.Image,
@@ -145,6 +162,9 @@ 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)
     
@@ -159,17 +179,12 @@ def infer_lighting_edit(
     
     prompt = build_lighting_prompt(azimuth, elevation)
     print(f"Generated Prompt: {prompt}")
-    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-        
     generator = torch.Generator(device=device).manual_seed(seed)
-    
     if image is None:
         raise gr.Error("Please upload an image first.")
-        
     pil_image = image.convert("RGB") if isinstance(image, Image.Image) else Image.open(image).convert("RGB")
-    
     result = pipe(
         image=[pil_image],
         prompt=prompt,
@@ -180,10 +195,10 @@ def infer_lighting_edit(
         guidance_scale=guidance_scale,
         num_images_per_prompt=1,
     ).images[0]
-    
     return result, seed, prompt
 
 def update_dimensions_on_upload(image):
+    """Compute recommended dimensions preserving aspect ratio."""
     if image is None:
         return 1024, 1024
     original_width, original_height = image.size
@@ -199,18 +214,19 @@ def update_dimensions_on_upload(image):
     new_height = (new_height // 8) * 8
     return new_width, new_height
 
-# --- 3D Control Component ---
 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}
         
         html_template = """
-        <div id="lighting-control-wrapper" style="width: 100%; height: 450px; position: relative; background: #1a1a1a; border-radius: 12px; overflow: hidden; border: 1px solid #333;">
-            <div id="prompt-overlay" style="position: absolute; bottom: 15px; left: 50%; transform: translateX(-50%); background: rgba(0,0,0,0.85); padding: 8px 16px; border-radius: 8px; font-family: 'IBM Plex Mono', monospace; font-size: 13px; color: #00ff88; white-space: nowrap; z-index: 10; border: 1px solid #333;"></div>
+        <div id="lighting-control-wrapper" style="width: 100%; height: 450px; position: relative; background: #1a1a1a; border-radius: 12px; overflow: hidden;">
+            <div id="prompt-overlay" style="position: absolute; bottom: 10px; left: 50%; transform: translateX(-50%); background: rgba(0,0,0,0.8); padding: 8px 16px; border-radius: 8px; font-family: monospace; font-size: 12px; color: #00ff88; white-space: nowrap; z-index: 10;"></div>
         </div>
         """
         
@@ -219,100 +235,110 @@ 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 ---
+                // Scene setup
                 const scene = new THREE.Scene();
                 scene.background = new THREE.Color(0x1a1a1a);
                 
                 const camera = new THREE.PerspectiveCamera(50, wrapper.clientWidth / wrapper.clientHeight, 0.1, 1000);
-                camera.position.set(4.5, 3.5, 4.5);
+                camera.position.set(4.5, 3, 4.5);
                 camera.lookAt(0, 0.75, 0);
                 
-                const renderer = new THREE.WebGLRenderer({ antialias: true, alpha: true });
+                const renderer = new THREE.WebGLRenderer({ antialias: true });
                 renderer.setSize(wrapper.clientWidth, wrapper.clientHeight);
                 renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
                 renderer.shadowMap.enabled = true;
                 renderer.shadowMap.type = THREE.PCFSoftShadowMap;
                 wrapper.insertBefore(renderer.domElement, promptOverlay);
                 
-                // --- Environment ---
-                scene.add(new THREE.AmbientLight(0xffffff, 0.15));
+                // Lighting
+                scene.add(new THREE.AmbientLight(0xffffff, 0.1));
                 
-                // Ground Plane
+                // Ground plane for shadows
                 const ground = new THREE.Mesh(
-                    new THREE.PlaneGeometry(12, 12),
-                    new THREE.ShadowMaterial({ opacity: 0.25 })
+                    new THREE.PlaneGeometry(10, 10),
+                    new THREE.ShadowMaterial({ opacity: 0.3 })
                 );
                 ground.rotation.x = -Math.PI / 2;
                 ground.position.y = 0;
                 ground.receiveShadow = true;
                 scene.add(ground);
                 
-                scene.add(new THREE.GridHelper(8, 16, 0x444444, 0x222222));
+                // Grid
+                scene.add(new THREE.GridHelper(8, 16, 0x333333, 0x222222));
                 
                 // Constants
                 const CENTER = new THREE.Vector3(0, 0.75, 0);
-                const BASE_DISTANCE = 2.6;
+                const BASE_DISTANCE = 2.5;
                 const AZIMUTH_RADIUS = 2.4;
-                const ELEVATION_RADIUS = 1.9;
-                
-                // --- Handle Colors (User Request) ---
-                const COLOR_AZIMUTH = 0xffd700; // Yellow/Gold
-                const COLOR_ELEVATION = 0x007bff; // Blue
+                const ELEVATION_RADIUS = 1.8;
                 
+                // State
                 let azimuthAngle = props.value?.azimuth || 0;
                 let elevationAngle = props.value?.elevation || 0;
                 
-                // --- Target Image Plane ---
+                // Mappings
+                const azimuthSteps = [0, 45, 90, 135, 180, 225, 270, 315];
+                const elevationSteps = [-90, 0, 90];
+                const azimuthNames = {
+                    0: 'Front', 45: 'Right Front', 90: 'Right',
+                    135: 'Right Rear', 180: 'Rear', 225: 'Left Rear',
+                    270: 'Left', 315: 'Left Front'
+                };
+                const elevationNames = { '-90': 'Below', '0': '', '90': 'Above' };
+                
+                function snapToNearest(value, steps) {
+                    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 = 512; canvas.height = 512;
+                    canvas.width = 256;
+                    canvas.height = 256;
                     const ctx = canvas.getContext('2d');
-                    // Gradient background
-                    const grd = ctx.createLinearGradient(0, 0, 0, 512);
-                    grd.addColorStop(0, '#2a2a3a');
-                    grd.addColorStop(1, '#1a1a2a');
-                    ctx.fillStyle = grd;
-                    ctx.fillRect(0, 0, 512, 512);
-                    
-                    // Simple silhouette
-                    ctx.fillStyle = '#ff4500';
+                    ctx.fillStyle = '#3a3a4a';
+                    ctx.fillRect(0, 0, 256, 256);
+                    ctx.fillStyle = '#ffcc99';
                     ctx.beginPath();
-                    ctx.arc(256, 256, 150, 0, Math.PI * 2);
+                    ctx.arc(128, 128, 80, 0, Math.PI * 2);
                     ctx.fill();
-                    
-                    ctx.fillStyle = '#ffffff';
-                    ctx.font = 'bold 40px sans-serif';
-                    ctx.textAlign = 'center';
-                    ctx.fillText("Preview", 256, 270);
-                    
+                    ctx.fillStyle = '#333';
+                    ctx.beginPath();
+                    ctx.arc(100, 110, 10, 0, Math.PI * 2);
+                    ctx.arc(156, 110, 10, 0, Math.PI * 2);
+                    ctx.fill();
+                    ctx.strokeStyle = '#333';
+                    ctx.lineWidth = 3;
+                    ctx.beginPath();
+                    ctx.arc(128, 130, 35, 0.2, Math.PI - 0.2);
+                    ctx.stroke();
                     return new THREE.CanvasTexture(canvas);
                 }
                 
+                // Target image plane
                 let currentTexture = createPlaceholderTexture();
-                const planeMaterial = new THREE.MeshStandardMaterial({ 
-                    map: currentTexture, 
-                    side: THREE.DoubleSide, 
-                    roughness: 0.4, 
-                    metalness: 0.1 
-                });
-                
-                let targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.5, 1.5), planeMaterial);
+                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);
                 targetPlane.position.copy(CENTER);
                 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.5, 1.5), planeMaterial);
+                        targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.2, 1.2), planeMaterial);
                         targetPlane.position.copy(CENTER);
                         targetPlane.receiveShadow = true;
                         scene.add(targetPlane);
@@ -324,71 +350,82 @@ class LightingControl3D(gr.HTML):
                     loader.load(url, (texture) => {
                         texture.minFilter = THREE.LinearFilter;
                         texture.magFilter = THREE.LinearFilter;
-                        texture.colorSpace = THREE.SRGBColorSpace;
                         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;
-                            const maxSize = 1.8;
-                            let w, h;
-                            if (aspect > 1) { w = maxSize; h = maxSize / aspect; } 
-                            else { h = maxSize; w = maxSize * aspect; }
-                            
+                            const maxSize = 1.5;
+                            let planeWidth, planeHeight;
+                            if (aspect > 1) {
+                                planeWidth = maxSize;
+                                planeHeight = maxSize / aspect;
+                            } else {
+                                planeHeight = maxSize;
+                                planeWidth = maxSize * aspect;
+                            }
                             scene.remove(targetPlane);
-                            targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(w, h), planeMaterial);
+                            targetPlane = new THREE.Mesh(
+                                new THREE.PlaneGeometry(planeWidth, planeHeight),
+                                planeMaterial
+                            );
                             targetPlane.position.copy(CENTER);
                             targetPlane.receiveShadow = true;
                             scene.add(targetPlane);
                         }
+                    }, undefined, (err) => {
+                        console.error('Failed to load texture:', err);
                     });
                 }
                 
-                if (props.imageUrl) updateTextureFromUrl(props.imageUrl);
+                // Check for initial imageUrl
+                if (props.imageUrl) {
+                    updateTextureFromUrl(props.imageUrl);
+                }
                 
-                // --- NEW LIGHT MODEL (Softbox Design) ---
-                const lightGroup = new THREE.Group();
+                // Create LED texture
+                function createLEDTexture() {
+                    const canvas = document.createElement('canvas');
+                    canvas.width = 64;
+                    canvas.height = 64;
+                    const ctx = canvas.getContext('2d');
+                    ctx.fillStyle = '#000000';
+                    ctx.fillRect(0, 0, 64, 64);
+                    for (let i = 0; i < 8; i++) {
+                        for (let j = 0; j < 8; j++) {
+                            ctx.fillStyle = '#ffffff';
+                            ctx.beginPath();
+                            ctx.arc(4 + 8 * i, 4 + 8 * j, 3, 0, Math.PI * 2);
+                            ctx.fill();
+                        }
+                    }
+                    return new THREE.CanvasTexture(canvas);
+                }
                 
-                // 1. Softbox Housing (Black matte box)
-                const housingGeo = new THREE.BoxGeometry(0.5, 0.5, 0.6); // Deep box
-                const housingMat = new THREE.MeshStandardMaterial({ 
-                    color: 0x222222, 
-                    roughness: 0.8, 
-                    metalness: 0.2 
-                });
-                const housing = new THREE.Mesh(housingGeo, housingMat);
-                // Shift housing back so origin is near front
-                housing.position.z = -0.15; 
-                lightGroup.add(housing);
-                
-                // 2. Connector / Bracket (Cylinder at back)
-                const bracketGeo = new THREE.CylinderGeometry(0.08, 0.08, 0.3, 16);
-                const bracketMat = new THREE.MeshStandardMaterial({ color: 0x555555, roughness: 0.5, metalness: 0.8 });
-                const bracket = new THREE.Mesh(bracketGeo, bracketMat);
-                bracket.rotation.x = Math.PI / 2;
-                bracket.position.z = -0.55;
-                lightGroup.add(bracket);
-
-                // 3. Diffuser Face (Bright White Emissive Plane)
-                const diffuserGeo = new THREE.PlaneGeometry(0.45, 0.45);
-                const diffuserMat = new THREE.MeshStandardMaterial({ 
-                    color: 0xffffff, 
-                    emissive: 0xffffff, 
-                    emissiveIntensity: 1.5,
-                    side: THREE.FrontSide
+                // Studio LED light model
+                const lightGroup = new THREE.Group();
+                const ledTexture = createLEDTexture();
+                const ledMat = new THREE.MeshStandardMaterial({
+                    color: 0x111111,
+                    emissive: 0xffffff,
+                    emissiveIntensity: 2,
+                    emissiveMap: ledTexture,
+                    roughness: 0.5,
+                    metalness: 0
                 });
-                const diffuser = new THREE.Mesh(diffuserGeo, diffuserMat);
-                diffuser.position.z = 0.151; // Just in front of housing front edge
-                lightGroup.add(diffuser);
+                const ledPanel = new THREE.Mesh(new THREE.BoxGeometry(0.6, 0.6, 0.1), ledMat);
+                lightGroup.add(ledPanel);
                 
-                // 4. Actual SpotLight Source
-                const spotLight = new THREE.SpotLight(0xffffff, 15, 12, Math.PI / 4, 0.5, 1);
-                spotLight.position.set(0, 0, 0.2); // Positioned at the face
+                const spotLight = new THREE.SpotLight(0xffffff, 10, 10, Math.PI / 3, 1, 1);
+                spotLight.position.set(0, 0, -0.05);
                 spotLight.castShadow = true;
                 spotLight.shadow.mapSize.width = 1024;
                 spotLight.shadow.mapSize.height = 1024;
-                spotLight.shadow.bias = -0.001;
+                spotLight.shadow.camera.near = 0.5;
+                spotLight.shadow.camera.far = 500;
+                spotLight.shadow.bias = -0.005;
                 lightGroup.add(spotLight);
                 
                 const lightTarget = new THREE.Object3D();
@@ -398,91 +435,62 @@ class LightingControl3D(gr.HTML):
                 
                 scene.add(lightGroup);
                 
-                // --- CONTROLS: Azimuth (Yellow) ---
+                // YELLOW: Azimuth ring
                 const azimuthRing = new THREE.Mesh(
-                    new THREE.TorusGeometry(AZIMUTH_RADIUS, 0.03, 16, 100),
-                    new THREE.MeshStandardMaterial({ 
-                        color: COLOR_AZIMUTH, 
-                        emissive: COLOR_AZIMUTH, 
-                        emissiveIntensity: 0.2,
-                        transparent: true,
-                        opacity: 0.6
-                    })
+                    new THREE.TorusGeometry(AZIMUTH_RADIUS, 0.04, 16, 64),
+                    new THREE.MeshStandardMaterial({ color: 0xffcc00, emissive: 0xffcc00, emissiveIntensity: 0.3 })
                 );
                 azimuthRing.rotation.x = Math.PI / 2;
                 azimuthRing.position.y = 0.05;
                 scene.add(azimuthRing);
                 
                 const azimuthHandle = new THREE.Mesh(
-                    new THREE.SphereGeometry(0.15, 32, 32),
-                    new THREE.MeshStandardMaterial({ 
-                        color: COLOR_AZIMUTH, 
-                        emissive: COLOR_AZIMUTH, 
-                        emissiveIntensity: 0.6,
-                        metalness: 0.5,
-                        roughness: 0.2
-                    })
+                    new THREE.SphereGeometry(0.18, 16, 16),
+                    new THREE.MeshStandardMaterial({ color: 0xffcc00, emissive: 0xffcc00, emissiveIntensity: 0.5 })
                 );
                 azimuthHandle.userData.type = 'azimuth';
                 scene.add(azimuthHandle);
                 
-                // --- CONTROLS: Elevation (Blue) ---
+                // BLUE: Elevation arc
                 const arcPoints = [];
-                for (let i = 0; i <= 40; i++) {
-                    const angle = THREE.MathUtils.degToRad(-90 + (180 * i / 40));
-                    arcPoints.push(new THREE.Vector3(-0.9, ELEVATION_RADIUS * Math.sin(angle) + CENTER.y, ELEVATION_RADIUS * Math.cos(angle)));
+                for (let i = 0; i <= 32; i++) {
+                    const angle = THREE.MathUtils.degToRad(-90 + (180 * i / 32));
+                    arcPoints.push(new THREE.Vector3(-0.8, ELEVATION_RADIUS * Math.sin(angle) + CENTER.y, ELEVATION_RADIUS * Math.cos(angle)));
                 }
                 const arcCurve = new THREE.CatmullRomCurve3(arcPoints);
                 const elevationArc = new THREE.Mesh(
-                    new THREE.TubeGeometry(arcCurve, 40, 0.03, 8, false),
-                    new THREE.MeshStandardMaterial({ 
-                        color: COLOR_ELEVATION, 
-                        emissive: COLOR_ELEVATION, 
-                        emissiveIntensity: 0.2,
-                        transparent: true,
-                        opacity: 0.6
-                    })
+                    new THREE.TubeGeometry(arcCurve, 32, 0.04, 8, false),
+                    new THREE.MeshStandardMaterial({ color: 0x007bff, emissive: 0x007bff, emissiveIntensity: 0.3 })
                 );
                 scene.add(elevationArc);
                 
                 const elevationHandle = new THREE.Mesh(
-                    new THREE.SphereGeometry(0.15, 32, 32),
-                    new THREE.MeshStandardMaterial({ 
-                        color: COLOR_ELEVATION, 
-                        emissive: COLOR_ELEVATION, 
-                        emissiveIntensity: 0.6,
-                        metalness: 0.5,
-                        roughness: 0.2
-                    })
+                    new THREE.SphereGeometry(0.18, 16, 16),
+                    new THREE.MeshStandardMaterial({ color: 0x007bff, emissive: 0x007bff, emissiveIntensity: 0.5 })
                 );
                 elevationHandle.userData.type = 'elevation';
                 scene.add(elevationHandle);
                 
-                // --- NEW REFRESH BUTTON ---
+                // REFRESH BUTTON
                 const refreshBtn = document.createElement('button');
-                refreshBtn.innerHTML = '↺ Reset View';
-                // Styling
-                Object.assign(refreshBtn.style, {
-                    position: 'absolute',
-                    top: '15px',
-                    right: '15px',
-                    background: 'linear-gradient(90deg, #ff5722, #d84315)',
-                    color: 'white',
-                    border: 'none',
-                    padding: '8px 16px',
-                    borderRadius: '6px',
-                    cursor: 'pointer',
-                    zIndex: '20',
-                    fontFamily: 'sans-serif',
-                    fontSize: '13px',
-                    fontWeight: '600',
-                    boxShadow: '0 4px 6px rgba(0,0,0,0.3)',
-                    transition: 'transform 0.1s ease, box-shadow 0.1s ease'
-                });
-                
-                refreshBtn.onmouseover = () => { refreshBtn.style.transform = 'translateY(-1px)'; refreshBtn.style.boxShadow = '0 6px 8px rgba(0,0,0,0.4)'; };
-                refreshBtn.onmouseout = () => { refreshBtn.style.transform = 'translateY(0)'; refreshBtn.style.boxShadow = '0 4px 6px rgba(0,0,0,0.3)'; };
-                refreshBtn.onmousedown = () => { refreshBtn.style.transform = 'translateY(1px)'; refreshBtn.style.boxShadow = '0 2px 4px rgba(0,0,0,0.3)'; };
+                refreshBtn.textContent = 'Reset View';
+                refreshBtn.style.position = 'absolute';
+                refreshBtn.style.top = '15px';
+                refreshBtn.style.right = '15px';
+                refreshBtn.style.backgroundColor = '#333';
+                refreshBtn.style.color = '#fff';
+                refreshBtn.style.border = '1px solid #555';
+                refreshBtn.style.borderRadius = '5px';
+                refreshBtn.style.padding = '5px 10px';
+                refreshBtn.style.cursor = 'pointer';
+                refreshBtn.style.zIndex = '20';
+                refreshBtn.style.fontSize = '12px';
+                refreshBtn.style.fontWeight = '500';
+                refreshBtn.style.transition = 'background-color 0.2s';
+                
+                // Hover effects
+                refreshBtn.onmouseover = function() { this.style.backgroundColor = '#444'; };
+                refreshBtn.onmouseout = function() { this.style.backgroundColor = '#333'; };
                 
                 wrapper.appendChild(refreshBtn);
                 
@@ -493,20 +501,11 @@ class LightingControl3D(gr.HTML):
                     updatePropsAndTrigger();
                 });
                 
-                // --- Logic & Events ---
-                const azimuthSteps = [0, 45, 90, 135, 180, 225, 270, 315];
-                const elevationSteps = [-90, 0, 90];
-                
-                function snapToNearest(value, steps) {
-                    return steps.reduce((prev, curr) => Math.abs(curr - value) < Math.abs(prev - value) ? curr : prev);
-                }
-                
                 function updatePositions() {
                     const distance = BASE_DISTANCE;
                     const azRad = THREE.MathUtils.degToRad(azimuthAngle);
                     const elRad = THREE.MathUtils.degToRad(elevationAngle);
                     
-                    // Position light group
                     const lightX = distance * Math.sin(azRad) * Math.cos(elRad);
                     const lightY = distance * Math.sin(elRad) + CENTER.y;
                     const lightZ = distance * Math.cos(azRad) * Math.cos(elRad);
@@ -514,181 +513,243 @@ class LightingControl3D(gr.HTML):
                     lightGroup.position.set(lightX, lightY, lightZ);
                     lightGroup.lookAt(CENTER);
                     
-                    // Position handles
-                    azimuthHandle.position.set(
-                        AZIMUTH_RADIUS * Math.sin(azRad), 
-                        0.05, 
-                        AZIMUTH_RADIUS * Math.cos(azRad)
-                    );
-                    
-                    elevationHandle.position.set(
-                        -0.9, 
-                        ELEVATION_RADIUS * Math.sin(elRad) + CENTER.y, 
-                        ELEVATION_RADIUS * Math.cos(elRad)
-                    );
+                    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 Text Prompt
+                    // Update prompt
                     const azSnap = snapToNearest(azimuthAngle, azimuthSteps);
                     const elSnap = snapToNearest(elevationAngle, elevationSteps);
-                    const azName = {0:'Front',45:'Right Front',90:'Right',135:'Right Rear',180:'Rear',225:'Left Rear',270:'Left',315:'Left Front'};
-                    const elName = {'-90':'Below','0':'','90':'Above'};
-                    
-                    let txt = 'Light source from';
-                    if (elSnap !== 0) txt += ' ' + elName[String(elSnap)];
-                    else txt += ' the ' + azName[azSnap];
-                    promptOverlay.textContent = txt;
+                    let prompt = 'Light source from';
+                    if (elSnap !== 0) {
+                        prompt += ' ' + elevationNames[String(elSnap)];
+                    } else {
+                        prompt += ' the ' + azimuthNames[azSnap];
+                    }
+                    promptOverlay.textContent = prompt;
                 }
                 
                 function updatePropsAndTrigger() {
                     const azSnap = snapToNearest(azimuthAngle, azimuthSteps);
                     const elSnap = snapToNearest(elevationAngle, elevationSteps);
+                    
                     props.value = { azimuth: azSnap, elevation: elSnap };
                     trigger('change', props.value);
                 }
                 
-                // Raycasting / Dragging
+                // Raycasting
                 const raycaster = new THREE.Raycaster();
                 const mouse = new THREE.Vector2();
                 let isDragging = false;
                 let dragTarget = null;
+                let dragStartMouse = new THREE.Vector2();
                 const intersection = new THREE.Vector3();
+                
                 const canvas = renderer.domElement;
                 
-                const updateMouse = (e) => {
+                canvas.addEventListener('mousedown', (e) => {
                     const rect = canvas.getBoundingClientRect();
-                    const clientX = e.touches ? e.touches[0].clientX : e.clientX;
-                    const clientY = e.touches ? e.touches[0].clientY : e.clientY;
-                    mouse.x = ((clientX - rect.left) / rect.width) * 2 - 1;
-                    mouse.y = -((clientY - rect.top) / rect.height) * 2 + 1;
-                };
-
-                const onStart = (e) => {
-                    e.preventDefault();
-                    updateMouse(e);
+                    mouse.x = ((e.clientX - rect.left) / rect.width) * 2 - 1;
+                    mouse.y = -((e.clientY - rect.top) / rect.height) * 2 + 1;
+                    
                     raycaster.setFromCamera(mouse, camera);
                     const intersects = raycaster.intersectObjects([azimuthHandle, elevationHandle]);
+                    
                     if (intersects.length > 0) {
                         isDragging = true;
                         dragTarget = intersects[0].object;
-                        dragTarget.scale.setScalar(1.3);
                         dragTarget.material.emissiveIntensity = 1.0;
+                        dragTarget.scale.setScalar(1.3);
+                        dragStartMouse.copy(mouse);
                         canvas.style.cursor = 'grabbing';
                     }
-                };
-
-                const onMove = (e) => {
-                    updateMouse(e);
+                });
+                
+                canvas.addEventListener('mousemove', (e) => {
+                    const rect = canvas.getBoundingClientRect();
+                    mouse.x = ((e.clientX - rect.left) / rect.width) * 2 - 1;
+                    mouse.y = -((e.clientY - rect.top) / rect.height) * 2 + 1;
                     
                     if (isDragging && dragTarget) {
                         raycaster.setFromCamera(mouse, camera);
+                        
                         if (dragTarget.userData.type === 'azimuth') {
                             const plane = new THREE.Plane(new THREE.Vector3(0, 1, 0), -0.05);
                             if (raycaster.ray.intersectPlane(plane, intersection)) {
                                 azimuthAngle = THREE.MathUtils.radToDeg(Math.atan2(intersection.x, intersection.z));
                                 if (azimuthAngle < 0) azimuthAngle += 360;
                             }
-                        } else {
-                            const plane = new THREE.Plane(new THREE.Vector3(1, 0, 0), -0.9);
+                        } else if (dragTarget.userData.type === 'elevation') {
+                            const plane = new THREE.Plane(new THREE.Vector3(1, 0, 0), -0.8);
                             if (raycaster.ray.intersectPlane(plane, intersection)) {
                                 const relY = intersection.y - CENTER.y;
-                                elevationAngle = THREE.MathUtils.clamp(
-                                    THREE.MathUtils.radToDeg(Math.atan2(relY, intersection.z)), 
-                                    -90, 90
-                                );
+                                const relZ = intersection.z;
+                                elevationAngle = THREE.MathUtils.clamp(THREE.MathUtils.radToDeg(Math.atan2(relY, relZ)), -90, 90);
                             }
                         }
                         updatePositions();
                     } else {
-                        // Hover effect
                         raycaster.setFromCamera(mouse, camera);
                         const intersects = raycaster.intersectObjects([azimuthHandle, elevationHandle]);
-                        canvas.style.cursor = intersects.length > 0 ? 'grab' : 'default';
                         [azimuthHandle, elevationHandle].forEach(h => {
-                            if (h !== dragTarget) {
-                                h.scale.setScalar(1);
-                                h.material.emissiveIntensity = 0.6;
-                            }
+                            h.material.emissiveIntensity = 0.5;
+                            h.scale.setScalar(1);
                         });
                         if (intersects.length > 0) {
-                            intersects[0].object.scale.setScalar(1.2);
-                            intersects[0].object.material.emissiveIntensity = 1.0;
+                            intersects[0].object.material.emissiveIntensity = 0.8;
+                            intersects[0].object.scale.setScalar(1.1);
+                            canvas.style.cursor = 'grab';
+                        } else {
+                            canvas.style.cursor = 'default';
                         }
                     }
-                };
+                });
                 
-                const onEnd = () => {
+                const onMouseUp = () => {
                     if (dragTarget) {
+                        dragTarget.material.emissiveIntensity = 0.5;
                         dragTarget.scale.setScalar(1);
-                        // Snap animation
+                        
+                        // Snap and animate
                         const targetAz = snapToNearest(azimuthAngle, azimuthSteps);
                         const targetEl = snapToNearest(elevationAngle, elevationSteps);
-                        const startAz = azimuthAngle;
-                        const startEl = elevationAngle;
+                        
+                        const startAz = azimuthAngle, startEl = elevationAngle;
                         const startTime = Date.now();
                         
-                        function anim() {
-                            const t = Math.min((Date.now() - startTime) / 250, 1);
+                        function animateSnap() {
+                            const t = Math.min((Date.now() - startTime) / 200, 1);
                             const ease = 1 - Math.pow(1 - t, 3);
                             
-                            // Azimuth wrapping logic
                             let azDiff = targetAz - startAz;
                             if (azDiff > 180) azDiff -= 360;
                             if (azDiff < -180) azDiff += 360;
-                            
                             azimuthAngle = startAz + azDiff * ease;
                             if (azimuthAngle < 0) azimuthAngle += 360;
                             if (azimuthAngle >= 360) azimuthAngle -= 360;
                             
                             elevationAngle = startEl + (targetEl - startEl) * ease;
-                            updatePositions();
                             
-                            if (t < 1) requestAnimationFrame(anim);
+                            updatePositions();
+                            if (t < 1) requestAnimationFrame(animateSnap);
                             else updatePropsAndTrigger();
                         }
-                        anim();
+                        animateSnap();
                     }
                     isDragging = false;
                     dragTarget = null;
                     canvas.style.cursor = 'default';
                 };
                 
-                canvas.addEventListener('mousedown', onStart);
-                window.addEventListener('mousemove', onMove);
-                window.addEventListener('mouseup', onEnd);
+                canvas.addEventListener('mouseup', onMouseUp);
+                canvas.addEventListener('mouseleave', onMouseUp);
+                # Touch support for mobile
+                canvas.addEventListener('touchstart', (e) => {
+                    e.preventDefault();
+                    const touch = e.touches[0];
+                    const rect = canvas.getBoundingClientRect();
+                    mouse.x = ((touch.clientX - rect.left) / rect.width) * 2 - 1;
+                    mouse.y = -((touch.clientY - rect.top) / rect.height) * 2 + 1;
+                    
+                    raycaster.setFromCamera(mouse, camera);
+                    const intersects = raycaster.intersectObjects([azimuthHandle, elevationHandle]);
+                    
+                    if (intersects.length > 0) {
+                        isDragging = true;
+                        dragTarget = intersects[0].object;
+                        dragTarget.material.emissiveIntensity = 1.0;
+                        dragTarget.scale.setScalar(1.3);
+                        dragStartMouse.copy(mouse);
+                    }
+                }, { passive: false });
+                
+                canvas.addEventListener('touchmove', (e) => {
+                    e.preventDefault();
+                    const touch = e.touches[0];
+                    const rect = canvas.getBoundingClientRect();
+                    mouse.x = ((touch.clientX - rect.left) / rect.width) * 2 - 1;
+                    mouse.y = -((touch.clientY - rect.top) / rect.height) * 2 + 1;
+                    
+                    if (isDragging && dragTarget) {
+                        raycaster.setFromCamera(mouse, camera);
+                        
+                        if (dragTarget.userData.type === 'azimuth') {
+                            const plane = new THREE.Plane(new THREE.Vector3(0, 1, 0), -0.05);
+                            if (raycaster.ray.intersectPlane(plane, intersection)) {
+                                azimuthAngle = THREE.MathUtils.radToDeg(Math.atan2(intersection.x, intersection.z));
+                                if (azimuthAngle < 0) azimuthAngle += 360;
+                            }
+                        } else if (dragTarget.userData.type === 'elevation') {
+                            const plane = new THREE.Plane(new THREE.Vector3(1, 0, 0), -0.8);
+                            if (raycaster.ray.intersectPlane(plane, intersection)) {
+                                const relY = intersection.y - CENTER.y;
+                                const relZ = intersection.z;
+                                elevationAngle = THREE.MathUtils.clamp(THREE.MathUtils.radToDeg(Math.atan2(relY, relZ)), -90, 90);
+                            }
+                        }
+                        updatePositions();
+                    }
+                }, { passive: false });
+                
+                canvas.addEventListener('touchend', (e) => {
+                    e.preventDefault();
+                    onMouseUp();
+                }, { passive: false });
                 
-                canvas.addEventListener('touchstart', onStart, {passive: false});
-                window.addEventListener('touchmove', onMove, {passive: false});
-                window.addEventListener('touchend', onEnd);
+                canvas.addEventListener('touchcancel', (e) => {
+                    e.preventDefault();
+                    onMouseUp();
+                }, { passive: false });
                 
-                // Initialization
+                // Initial update
                 updatePositions();
                 
-                const animate = () => {
-                    requestAnimationFrame(animate);
+                // Render loop
+                function render() {
+                    requestAnimationFrame(render);
                     renderer.render(scene, camera);
-                };
-                animate();
+                }
+                render();
                 
-                // Props Watcher
-                setInterval(() => {
-                    if (props.imageUrl !== currentTexture.sourceFile) { 
-                        // simple check, better logic exists but keeping simple
-                    }
-                    if (JSON.stringify(props.value) !== JSON.stringify({azimuth:snapToNearest(azimuthAngle, azimuthSteps), elevation:snapToNearest(elevationAngle, elevationSteps)})) {
-                         // External update logic here if needed
-                    }
-                }, 200);
+                // Handle resize
+                new ResizeObserver(() => {
+                    camera.aspect = wrapper.clientWidth / wrapper.clientHeight;
+                    camera.updateProjectionMatrix();
+                    renderer.setSize(wrapper.clientWidth, wrapper.clientHeight);
+                }).observe(wrapper);
                 
-                // External Interface
+                // Store update functions for external calls
                 wrapper._updateFromProps = (newVal) => {
-                    if (newVal) {
+                    if (newVal && typeof newVal === 'object') {
                         azimuthAngle = newVal.azimuth ?? azimuthAngle;
                         elevationAngle = newVal.elevation ?? elevationAngle;
                         updatePositions();
                     }
                 };
+                
                 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;
+                        if (props.value && typeof props.value === 'object') {
+                            azimuthAngle = props.value.azimuth ?? azimuthAngle;
+                            elevationAngle = props.value.elevation ?? elevationAngle;
+                            updatePositions();
+                        }
+                    }
+                }, 100);
             };
+            
             initScene();
         })();
         """
@@ -701,7 +762,6 @@ class LightingControl3D(gr.HTML):
             **kwargs
         )
 
-# --- Gradio UI Layout ---
 css = '''
 #col-container { max-width: 1200px; margin: 0 auto; }
 .dark .progress-text { color: white !important; }
@@ -709,17 +769,16 @@ css = '''
 .slider-row { display: flex; gap: 10px; align-items: center; }
 #main-title h1 {font-size: 2.4em !important;}
 '''
-
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# **Qwen-Image-Edit-2511-3D-Lighting-Control**", elem_id="main-title")
-    gr.Markdown("Control lighting directions using the **3D viewport** or **sliders**. Use [Multi-Angle-Lighting](https://huggingface.co/dx8152/Qwen-Edit-2509-Multi-Angle-Lighting) LoRA for precise lighting control.")
+    gr.Markdown("Control lighting directions using the **3D viewport** or **sliders**. Use [Multi-Angle-Lighting](https://huggingface.co/dx8152/Qwen-Edit-2509-Multi-Angle-Lighting) LoRA for precise lighting control, paired with [Rapid-AIO-V19](https://huggingface.co/prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19).")
     
     with gr.Row():
         with gr.Column(scale=1):
             image = gr.Image(label="Input Image", type="pil", height=300)
             
             gr.Markdown("### 3D Lighting Control")
-            gr.Markdown("*Drag the colored handles: 🟡 **Yellow** Azimuth (Direction), 🔵 **Blue** Elevation (Height)*")
+            gr.Markdown("*Drag the colored handles: 🟡 Azimuth (Direction), 🔵 Elevation (Height)*")
             
             lighting_3d = LightingControl3D(
                 value={"azimuth": 0, "elevation": 0},
@@ -731,22 +790,29 @@ with gr.Blocks(css=css) as demo:
             
             azimuth_slider = gr.Slider(
                 label="Azimuth (Horizontal Rotation)",
-                minimum=0, maximum=315, step=45, value=0,
+                minimum=0,
+                maximum=315,
+                step=45,
+                value=0,
                 info="0°=front, 90°=right, 180°=rear, 270°=left"
             )
             
             elevation_slider = gr.Slider(
                 label="Elevation (Vertical Angle)",
-                minimum=-90, maximum=90, step=90, value=0,
+                minimum=-90,
+                maximum=90,
+                step=90,
+                value=0,
                 info="-90°=from below, 0°=horizontal, 90°=from above"
             )
 
-            prompt_preview = gr.Textbox(
-                label="Generated Prompt",
-                value="Light source from the Front",
-                interactive=True,
-                lines=1,
-            )
+            with gr.Row():
+                prompt_preview = gr.Textbox(
+                    label="Generated Prompt",
+                    value="Light source from the Front",
+                    interactive=True,
+                    lines=1,
+                )
         
         with gr.Column(scale=1):
             result = gr.Image(label="Output Image", height=500)
@@ -759,11 +825,13 @@ with gr.Blocks(css=css) as demo:
                 height = gr.Slider(label="Height", minimum=256, maximum=2048, step=8, value=1024)
                 width = gr.Slider(label="Width", minimum=256, maximum=2048, step=8, value=1024)
     
-    # --- Event Wiring ---
     def update_prompt_from_sliders(azimuth, elevation):
-        return build_lighting_prompt(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)
@@ -772,11 +840,14 @@ 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)
+
         import base64
         from io import BytesIO
         buffered = BytesIO()
@@ -825,7 +896,8 @@ with gr.Blocks(css=css) as demo:
         fn=lambda: gr.update(imageUrl=None),
         outputs=[lighting_3d]
     )
-
+    
 if __name__ == "__main__":
     head = '<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>'
-    demo.launch(head=head, theme=orange_red_theme, mcp_server=True, ssr_mode=False, show_error=True)
+    css = '.fillable{max-width: 1200px !important}'
+    demo.launch(head=head, css=css, theme=orange_red_theme, mcp_server=True, ssr_mode=False, show_error=True)