Update app.py

app.py

@@ -4,15 +4,26 @@ import random
 import torch
 import spaces
 from PIL import Image
+
+# --- Updated Imports ---
 from diffusers import FlowMatchEulerDiscreteScheduler
-from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
-from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
-from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+try:
+    from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
+    from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+    from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+except ImportError:
+    # Fallback for standard environments or if package missing (user should install)
+    print("Warning: 'qwenimage' package not found. Ensure it is installed.")
 
 MAX_SEED = np.iinfo(np.int32).max
-# --- Model Loading ---
+
+# --- Model Configuration ---
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print(f"Loading Model on {device}...")
+
+# Load the pipeline with the custom transformer
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2511",
     transformer=QwenImageTransformer2DModel.from_pretrained(
@@ -22,63 +33,53 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
     ),
     torch_dtype=dtype
 ).to(device)
+
+# Attempt to set Flash Attention 3
 try:
     pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
     print("Flash Attention 3 Processor set successfully.")
 except Exception as e:
-    print(f"Warning: Could not set FA3 processor: {e}")
+    print(f"Standard Attention enabled (FA3 skipped: {e})")
 
+# Adapter Specs
 ADAPTER_SPECS = {
     "Multi-Angle-Lighting": {
         "repo": "dx8152/Qwen-Edit-2509-Multi-Angle-Lighting",
         "weights": "多角度灯光-251116.safetensors",
         "adapter_name": "multi-angle-lighting"
-    },
-}
-loaded = False
-
-# --- Prompt Building ---
-# Azimuth mappings (8 positions)
-AZIMUTH_MAP = {
-    0: "Front",
-    45: "Right Front",
-    90: "Right",
-    135: "Right Rear",
-    180: "Rear",
-    225: "Left Rear",
-    270: "Left",
-    315: "Left Front"
+    }
 }
-# Elevation mappings (3 positions)
-ELEVATION_MAP = {
-    -90: "Below",
-    0: "",
-    90: "Above"
+
+# State to track loaded adapter
+CURRENT_LOADED_ADAPTER = None
+
+# --- Prompt Building & Logic ---
+
+# 0 degrees = Front
+LIGHTING_AZIMUTH_MAP = {
+    0: "Light source from the Front",
+    45: "Light source from the Right Front",
+    90: "Light source from the Right",
+    135: "Light source from the Right Rear",
+    180: "Light source from the Rear",
+    225: "Light source from the Left Rear",
+    270: "Light source from the Left",
+    315: "Light source from the Left Front"
 }
 
 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()))
-    
-    if elevation_snapped == 0:
-        return f"Light source from the {AZIMUTH_MAP[azimuth_snapped]}"
-    else:
-        return f"Light source from {ELEVATION_MAP[elevation_snapped]}"
+    # Check Vertical extremes
+    if elevation >= 60:
+        return "Light source from Above"
+    if elevation <= -60:
+        return "Light source from Below"
+        
+    # Snap Horizontal
+    azimuth_snapped = snap_to_nearest(azimuth, list(LIGHTING_AZIMUTH_MAP.keys()))
+    return LIGHTING_AZIMUTH_MAP[azimuth_snapped]
 
 @spaces.GPU
 def infer_lighting_edit(
@@ -92,29 +93,33 @@ 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)
+    global CURRENT_LOADED_ADAPTER
     
-    if not loaded:
+    # --- Lazy Loading Logic ---
+    spec = ADAPTER_SPECS["Multi-Angle-Lighting"]
+    if CURRENT_LOADED_ADAPTER != spec["adapter_name"]:
+        print(f"Lazy loading adapter: {spec['adapter_name']}...")
         pipe.load_lora_weights(
-            ADAPTER_SPECS["Multi-Angle-Lighting"]["repo"],
-            weight_name=ADAPTER_SPECS["Multi-Angle-Lighting"]["weights"],
-            adapter_name=ADAPTER_SPECS["Multi-Angle-Lighting"]["adapter_name"]
+            spec["repo"],
+            weight_name=spec["weights"],
+            adapter_name=spec["adapter_name"]
         )
-        pipe.set_adapters([ADAPTER_SPECS["Multi-Angle-Lighting"]["adapter_name"]], adapter_weights=[1.0])
-        loaded = True
+        pipe.set_adapters([spec["adapter_name"]], adapter_weights=[1.0])
+        CURRENT_LOADED_ADAPTER = spec["adapter_name"]
     
+    # --- Generation ---
     prompt = build_lighting_prompt(azimuth, elevation)
-    print(f"Generated Prompt: {prompt}")
+    print(f"Generated Lighting 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,
@@ -125,10 +130,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
@@ -145,455 +150,318 @@ def update_dimensions_on_upload(image):
     return new_width, new_height
 
 # --- 3D Lighting Control Component ---
-class LightingControl3D(gr.HTML):
+class LightControl3D(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.
+    Advanced 3D Lighting control component using Three.js + OrbitControls.
     """
     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;">
-            <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 id="light-control-wrapper" style="width: 100%; height: 500px; position: relative; background: radial-gradient(circle at center, #1a1a1a 0%, #000000 100%); border-radius: 12px; overflow: hidden; border: 1px solid #333; box-shadow: inset 0 0 20px #000;">
+            <div id="prompt-overlay" style="position: absolute; top: 15px; left: 50%; transform: translateX(-50%); background: rgba(0, 0, 0, 0.85); padding: 8px 20px; border-radius: 30px; font-family: 'Segoe UI', sans-serif; font-weight: 600; font-size: 14px; color: #FFD700; white-space: nowrap; z-index: 10; border: 1px solid #FFD700; box-shadow: 0 0 10px rgba(255, 215, 0, 0.3);">
+                Light Source: Front
+            </div>
+            <div style="position: absolute; bottom: 15px; left: 15px; color: #666; font-family: monospace; font-size: 11px; z-index: 10;">
+                LEFT CLICK: Move Light • RIGHT CLICK: Rotate View
+            </div>
         </div>
         """
         
         js_on_load = """
         (() => {
-            const wrapper = element.querySelector('#lighting-control-wrapper');
+            const wrapper = element.querySelector('#light-control-wrapper');
             const promptOverlay = element.querySelector('#prompt-overlay');
             
-            // Wait for THREE to load
             const initScene = () => {
-                if (typeof THREE === 'undefined') {
+                if (typeof THREE === 'undefined' || typeof THREE.OrbitControls === 'undefined') {
                     setTimeout(initScene, 100);
                     return;
                 }
                 
-                // Scene setup
+                // 1. Setup Scene
                 const scene = new THREE.Scene();
-                scene.background = new THREE.Color(0x1a1a1a);
+                // No background color, letting CSS gradient show through
                 
-                const camera = new THREE.PerspectiveCamera(50, wrapper.clientWidth / wrapper.clientHeight, 0.1, 1000);
-                camera.position.set(4.5, 3, 4.5);
-                camera.lookAt(0, 0.75, 0);
+                const camera = new THREE.PerspectiveCamera(45, wrapper.clientWidth / wrapper.clientHeight, 0.1, 1000);
+                camera.position.set(4, 3, 4);
                 
-                const renderer = new THREE.WebGLRenderer({ antialias: true });
-                renderer.setSize(wrapper.clientWidth / wrapper.clientHeight);
+                const renderer = new THREE.WebGLRenderer({ antialias: true, alpha: true });
+                renderer.setSize(wrapper.clientWidth, wrapper.clientHeight);
                 renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
+                renderer.shadowMap.enabled = true; // Enable Shadows
+                renderer.shadowMap.type = THREE.PCFSoftShadowMap;
                 wrapper.insertBefore(renderer.domElement, promptOverlay);
                 
-                // Lighting
-                scene.add(new THREE.AmbientLight(0xffffff, 0.2));
-                
-                // Grid
-                scene.add(new THREE.GridHelper(8, 16, 0x333333, 0x222222));
-                
-                // Constants
+                // 2. Controls (Orbit for Camera)
+                const controls = new THREE.OrbitControls(camera, renderer.domElement);
+                controls.enableDamping = true;
+                controls.dampingFactor = 0.05;
+                controls.maxDistance = 10;
+                controls.minDistance = 2;
+                controls.target.set(0, 0.75, 0);
+
+                // 3. Environment
                 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 = {
-                    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' };
+                const ORBIT_RADIUS = 2.2;
                 
-                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 = 256;
-                    canvas.height = 256;
-                    const ctx = canvas.getContext('2d');
-                    ctx.fillStyle = '#3a3a4a';
-                    ctx.fillRect(0, 0, 256, 256);
-                    ctx.fillStyle = '#ffcc99';
-                    ctx.beginPath();
-                    ctx.arc(128, 128, 80, 0, Math.PI * 2);
-                    ctx.fill();
-                    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);
-                }
+                // Floor (Shadow Receiver)
+                const planeGeo = new THREE.PlaneGeometry(10, 10);
+                const planeMat = new THREE.ShadowMaterial({ opacity: 0.5, color: 0x000000 });
+                const floor = new THREE.Mesh(planeGeo, planeMat);
+                floor.rotation.x = -Math.PI / 2;
+                floor.receiveShadow = true;
+                scene.add(floor);
                 
-                // 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);
-                targetPlane.position.copy(CENTER);
-                scene.add(targetPlane);
-                
-                // Function to update texture from image URL
+                // Grid
+                const grid = new THREE.GridHelper(10, 20, 0x444444, 0x111111);
+                scene.add(grid);
+                
+                // Dome Helper (Wireframe)
+                const domeGeo = new THREE.SphereGeometry(ORBIT_RADIUS, 16, 12, 0, Math.PI*2, 0, Math.PI/2);
+                const domeMat = new THREE.MeshBasicMaterial({ color: 0x333333, wireframe: true, transparent: true, opacity: 0.15 });
+                const dome = new THREE.Mesh(domeGeo, domeMat);
+                dome.position.copy(CENTER);
+                dome.position.y -= 0.75; // Adjust to base
+                scene.add(dome);
+
+                // 4. The Subject (Image Box)
+                let currentTexture = null;
+                const boxGeo = new THREE.BoxGeometry(1.2, 1.2, 0.1); // Thin box instead of plane
+                const boxMat = new THREE.MeshStandardMaterial({ color: 0x222222, roughness: 0.7, metalness: 0.1 });
+                let subjectMesh = new THREE.Mesh(boxGeo, boxMat);
+                subjectMesh.position.copy(CENTER);
+                subjectMesh.castShadow = true;
+                subjectMesh.receiveShadow = true;
+                scene.add(subjectMesh);
+                
+                // Hidden Reference Sphere (To cast better shadows if image is flat)
+                const refSphere = new THREE.Mesh(
+                    new THREE.SphereGeometry(0.5, 32, 32),
+                    new THREE.MeshBasicMaterial({ visible: false }) // Invisible but casts shadow? No, must be visible. 
+                );
+                // Actually, let's keep the box, it casts decent shadows.
+
                 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);
-                        scene.add(targetPlane);
+                        boxMat.map = null;
+                        boxMat.needsUpdate = true;
                         return;
                     }
-                    
-                    const loader = new THREE.TextureLoader();
-                    loader.crossOrigin = 'anonymous';
-                    loader.load(url, (texture) => {
-                        texture.minFilter = THREE.LinearFilter;
-                        texture.magFilter = THREE.LinearFilter;
-                        planeMaterial.map = texture;
-                        planeMaterial.needsUpdate = true;
-                        
-                        // Adjust plane aspect ratio to match image
-                        const img = texture.image;
-                        if (img && img.width && img.height) {
+                    new THREE.TextureLoader().load(url, (tex) => {
+                        boxMat.map = tex;
+                        boxMat.color.setHex(0xffffff); // Reset color if texture present
+                        boxMat.needsUpdate = true;
+                        const img = tex.image;
+                        if(img.width && img.height) {
                             const aspect = img.width / img.height;
-                            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(planeWidth, planeHeight),
-                                planeMaterial
+                            const scale = aspect > 1 ? 1 : aspect;
+                            // Adjust box dimensions
+                            const newGeo = new THREE.BoxGeometry(
+                                aspect > 1 ? 1.5 : 1.5 * aspect, 
+                                aspect > 1 ? 1.5 / aspect : 1.5, 
+                                0.05
                             );
-                            targetPlane.position.copy(CENTER);
-                            scene.add(targetPlane);
+                            subjectMesh.geometry.dispose();
+                            subjectMesh.geometry = newGeo;
                         }
-                    }, undefined, (err) => {
-                        console.error('Failed to load texture:', err);
                     });
                 }
-                
-                // Check for initial imageUrl
-                if (props.imageUrl) {
-                    updateTextureFromUrl(props.imageUrl);
-                }
-                
-                // Light model
+                if (props.imageUrl) updateTextureFromUrl(props.imageUrl);
+
+                // 5. Light Source (The Sun) & The Actual Light
                 const lightGroup = new THREE.Group();
-                const bulbMat = new THREE.MeshStandardMaterial({ color: 0xffff00, emissive: 0xffff00, emissiveIntensity: 1 });
-                const bulb = new THREE.Mesh(new THREE.SphereGeometry(0.15, 16, 16), bulbMat);
-                lightGroup.add(bulb);
-                const pointLight = new THREE.PointLight(0xffffff, 5, 0);
-                lightGroup.add(pointLight);
                 scene.add(lightGroup);
                 
-                // GREEN: Azimuth ring
-                const azimuthRing = new THREE.Mesh(
-                    new THREE.TorusGeometry(AZIMUTH_RADIUS, 0.04, 16, 64),
-                    new THREE.MeshStandardMaterial({ color: 0x00ff88, emissive: 0x00ff88, emissiveIntensity: 0.3 })
+                // The Visual Sun
+                const sunMesh = new THREE.Mesh(
+                    new THREE.SphereGeometry(0.2, 32, 32),
+                    new THREE.MeshBasicMaterial({ color: 0xFFD700 })
                 );
-                azimuthRing.rotation.x = Math.PI / 2;
-                azimuthRing.position.y = 0.05;
-                scene.add(azimuthRing);
+                lightGroup.add(sunMesh);
                 
-                const azimuthHandle = new THREE.Mesh(
-                    new THREE.SphereGeometry(0.18, 16, 16),
-                    new THREE.MeshStandardMaterial({ color: 0x00ff88, emissive: 0x00ff88, emissiveIntensity: 0.5 })
+                // Glow
+                const glowMesh = new THREE.Mesh(
+                    new THREE.SphereGeometry(0.35, 32, 32),
+                    new THREE.MeshBasicMaterial({ color: 0xFFD700, transparent: true, opacity: 0.3 })
                 );
-                azimuthHandle.userData.type = 'azimuth';
-                scene.add(azimuthHandle);
-                
-                // PINK: Elevation arc
-                const arcPoints = [];
-                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)));
+                sunMesh.add(glowMesh);
+                
+                // The Actual Light Caster
+                const spotLight = new THREE.SpotLight(0xffffff, 2);
+                spotLight.angle = Math.PI / 4;
+                spotLight.penumbra = 0.5;
+                spotLight.decay = 2;
+                spotLight.distance = 10;
+                spotLight.castShadow = true;
+                spotLight.shadow.mapSize.width = 1024;
+                spotLight.shadow.mapSize.height = 1024;
+                lightGroup.add(spotLight);
+                
+                // Ambient fill
+                scene.add(new THREE.AmbientLight(0xffffff, 0.15));
+
+                // Target Line
+                const lineGeo = new THREE.BufferGeometry().setFromPoints([new THREE.Vector3(0,0,0), new THREE.Vector3(0,0,-1)]);
+                const lineMat = new THREE.LineBasicMaterial({ color: 0xFFD700, transparent: true, opacity: 0.5 });
+                const line = new THREE.Line(lineGeo, lineMat);
+                lightGroup.add(line);
+
+                // State
+                let azimuthAngle = props.value?.azimuth || 0;
+                let elevationAngle = props.value?.elevation || 0;
+
+                const azimuthNames = {
+                    0: 'Front', 45: 'Right Front', 90: 'Right',
+                    135: 'Right Rear', 180: 'Rear', 225: 'Left Rear',
+                    270: 'Left', 315: 'Left Front'
+                };
+
+                function getPromptText(az, el) {
+                    if (el >= 60) return "Light source from Above";
+                    if (el <= -60) return "Light source from Below";
+                    const steps = [0, 45, 90, 135, 180, 225, 270, 315];
+                    const snapped = steps.reduce((prev, curr) => Math.abs(curr - az) < Math.abs(prev - az) ? curr : prev);
+                    return "Light source from the " + azimuthNames[snapped];
                 }
-                const arcCurve = new THREE.CatmullRomCurve3(arcPoints);
-                const elevationArc = new THREE.Mesh(
-                    new THREE.TubeGeometry(arcCurve, 32, 0.04, 8, false),
-                    new THREE.MeshStandardMaterial({ color: 0xff69b4, emissive: 0xff69b4, emissiveIntensity: 0.3 })
-                );
-                scene.add(elevationArc);
-                
-                const elevationHandle = new THREE.Mesh(
-                    new THREE.SphereGeometry(0.18, 16, 16),
-                    new THREE.MeshStandardMaterial({ color: 0xff69b4, emissive: 0xff69b4, emissiveIntensity: 0.5 })
-                );
-                elevationHandle.userData.type = 'elevation';
-                scene.add(elevationHandle);
-                
+
                 function updatePositions() {
-                    const distance = BASE_DISTANCE;
                     const azRad = THREE.MathUtils.degToRad(azimuthAngle);
                     const elRad = THREE.MathUtils.degToRad(elevationAngle);
                     
-                    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);
+                    const x = ORBIT_RADIUS * Math.sin(azRad) * Math.cos(elRad);
+                    const y = ORBIT_RADIUS * Math.sin(elRad) + CENTER.y;
+                    const z = ORBIT_RADIUS * Math.cos(azRad) * Math.cos(elRad);
                     
-                    lightGroup.position.set(lightX, lightY, lightZ);
+                    lightGroup.position.set(x, y, z);
+                    lightGroup.lookAt(CENTER);
                     
-                    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';
-                    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);
+                    // Line length
+                    const dist = lightGroup.position.distanceTo(CENTER);
+                    line.scale.z = dist;
+
+                    // Update Prompt UI
+                    promptOverlay.textContent = getPromptText(azimuthAngle, elevationAngle);
                     
-                    props.value = { azimuth: azSnap, elevation: elSnap };
-                    trigger('change', props.value);
+                    // Color updates for warning zones
+                    const isExtreme = elevationAngle >= 60 || elevationAngle <= -60;
+                    const color = isExtreme ? 0xFF4500 : 0xFFD700;
+                    sunMesh.material.color.setHex(color);
+                    promptOverlay.style.color = isExtreme ? '#FF4500' : '#FFD700';
+                    promptOverlay.style.borderColor = isExtreme ? '#FF4500' : '#FFD700';
                 }
+
+                // --- Interaction Logic ---
+                // We need to distinguish between Orbiting (Background drag) and Light Moving (Sun drag)
                 
-                // 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;
+                // Drag Helper Sphere (Invisible large sphere for raycasting light position)
+                const dragSphere = new THREE.Mesh(
+                    new THREE.SphereGeometry(ORBIT_RADIUS, 32, 32),
+                    new THREE.MeshBasicMaterial({ visible: false, side: THREE.DoubleSide })
+                );
+                dragSphere.position.copy(CENTER);
+                scene.add(dragSphere);
                 
-                canvas.addEventListener('mousedown', (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;
+                let isDraggingLight = false;
+
+                function onDown(e) {
+                    const rect = wrapper.getBoundingClientRect();
+                    const x = e.clientX || e.touches[0].clientX;
+                    const y = e.clientY || e.touches[0].clientY;
+                    mouse.x = ((x - rect.left) / rect.width) * 2 - 1;
+                    mouse.y = -((y - 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);
-                        canvas.style.cursor = 'grabbing';
-                    }
-                });
-                
-                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;
+                    // Check if clicking the Sun
+                    const sunIntersects = raycaster.intersectObject(sunMesh);
+                    const glowIntersects = raycaster.intersectObject(glowMesh);
                     
-                    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();
-                    } else {
-                        raycaster.setFromCamera(mouse, camera);
-                        const intersects = raycaster.intersectObjects([azimuthHandle, elevationHandle]);
-                        [azimuthHandle, elevationHandle].forEach(h => {
-                            h.material.emissiveIntensity = 0.5;
-                            h.scale.setScalar(1);
-                        });
-                        if (intersects.length > 0) {
-                            intersects[0].object.material.emissiveIntensity = 0.8;
-                            intersects[0].object.scale.setScalar(1.1);
-                            canvas.style.cursor = 'grab';
-                        } else {
-                            canvas.style.cursor = 'default';
-                        }
+                    if (sunIntersects.length > 0 || glowIntersects.length > 0) {
+                        isDraggingLight = true;
+                        controls.enabled = false; // Disable Orbit
+                        wrapper.style.cursor = 'grabbing';
                     }
-                });
-                
-                const onMouseUp = () => {
-                    if (dragTarget) {
-                        dragTarget.material.emissiveIntensity = 0.5;
-                        dragTarget.scale.setScalar(1);
-                        
-                        // Snap and animate
-                        const targetAz = snapToNearest(azimuthAngle, azimuthSteps);
-                        const targetEl = snapToNearest(elevationAngle, elevationSteps);
-                        
-                        const startAz = azimuthAngle, startEl = elevationAngle;
-                        const startTime = Date.now();
-                        
-                        function animateSnap() {
-                            const t = Math.min((Date.now() - startTime) / 200, 1);
-                            const ease = 1 - Math.pow(1 - t, 3);
+                }
+
+                function onMove(e) {
+                    const rect = wrapper.getBoundingClientRect();
+                    const x = e.clientX || (e.touches ? e.touches[0].clientX : 0);
+                    const y = e.clientY || (e.touches ? e.touches[0].clientY : 0);
+                    mouse.x = ((x - rect.left) / rect.width) * 2 - 1;
+                    mouse.y = -((y - rect.top) / rect.height) * 2 + 1;
+                    
+                    raycaster.setFromCamera(mouse, camera);
+
+                    if (isDraggingLight) {
+                        const intersects = raycaster.intersectObject(dragSphere);
+                        if(intersects.length > 0) {
+                            const point = intersects[0].point;
+                            const rel = new THREE.Vector3().subVectors(point, CENTER);
                             
-                            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;
+                            let az = Math.atan2(rel.x, rel.z) * (180/Math.PI);
+                            if(az < 0) az += 360;
                             
-                            elevationAngle = startEl + (targetEl - startEl) * ease;
+                            const dist = Math.sqrt(rel.x*rel.x + rel.z*rel.z);
+                            let el = Math.atan2(rel.y, dist) * (180/Math.PI);
+                            el = Math.max(-89, Math.min(89, el)); // Clamp
                             
+                            azimuthAngle = az;
+                            elevationAngle = el;
                             updatePositions();
-                            if (t < 1) requestAnimationFrame(animateSnap);
-                            else updatePropsAndTrigger();
                         }
-                        animateSnap();
-                    }
-                    isDragging = false;
-                    dragTarget = null;
-                    canvas.style.cursor = 'default';
-                };
-                
-                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);
+                    } else {
+                        // Hover state
+                        const sunIntersects = raycaster.intersectObject(sunMesh);
+                        if (sunIntersects.length > 0) wrapper.style.cursor = 'grab';
+                        else wrapper.style.cursor = 'default';
                     }
-                }, { 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);
+                }
+
+                function onUp() {
+                    if (isDraggingLight) {
+                        isDraggingLight = false;
+                        controls.enabled = true; // Re-enable Orbit
+                        wrapper.style.cursor = 'default';
                         
-                        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();
+                        props.value = { azimuth: azimuthAngle, elevation: elevationAngle };
+                        trigger('change', props.value);
                     }
-                }, { passive: false });
-                
-                canvas.addEventListener('touchend', (e) => {
-                    e.preventDefault();
-                    onMouseUp();
-                }, { passive: false });
-                
-                canvas.addEventListener('touchcancel', (e) => {
-                    e.preventDefault();
-                    onMouseUp();
-                }, { passive: false });
-                
-                // Initial update
+                }
+
+                wrapper.addEventListener('mousedown', onDown);
+                window.addEventListener('mousemove', onMove);
+                window.addEventListener('mouseup', onUp);
+                wrapper.addEventListener('touchstart', onDown, {passive: false});
+                window.addEventListener('touchmove', onMove, {passive: false});
+                window.addEventListener('touchend', onUp);
+
                 updatePositions();
-                
-                // Render loop
-                function render() {
-                    requestAnimationFrame(render);
+
+                function animate() {
+                    requestAnimationFrame(animate);
+                    controls.update();
                     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;
-                        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);
+                animate();
+
+                // Prop watchers
                 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;
+                     if (props.value && (props.value.azimuth !== azimuthAngle || props.value.elevation !== elevationAngle)) {
+                        if(!isDraggingLight) {
+                            azimuthAngle = props.value.azimuth;
+                            elevationAngle = props.value.elevation;
                             updatePositions();
                         }
                     }
                 }, 100);
-            };
-            
+                
+                wrapper._updateTexture = updateTextureFromUrl;
+            }
             initScene();
         })();
         """
@@ -606,154 +474,95 @@ class LightingControl3D(gr.HTML):
             **kwargs
         )
 
-# --- UI ---
+# --- UI Layout ---
 css = '''
 #col-container { max-width: 1200px; margin: 0 auto; }
 .dark .progress-text { color: white !important; }
-#lighting-3d-control { min-height: 450px; }
-.slider-row { display: flex; gap: 10px; align-items: center; }
 '''
+
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # 🎬 Qwen Image Edit 2511 — 3D Lighting Control
+    # 💡 Qwen Edit 2509 — Studio Lighting Control
     
-    Control lighting directions using the **3D viewport** or **sliders**.
-    Using [dx8152/Qwen-Edit-2509-Multi-Angle-Lighting] for precise lighting control.
+    **Interactive 3D Studio:** * **Left Drag** the 🟡 **Sun** to change light direction.
+    * **Right Drag** background to rotate the camera view.
+    * The shadow on the floor previews the light direction.
     """)
     
     with gr.Row():
-        # Left column: Input image and controls
         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: 🟢 Azimuth (Direction), 🩷 Elevation (Height)*")
+            gr.Markdown("### 🎮 3D Studio Controller")
             
-            lighting_3d = LightingControl3D(
+            light_3d = LightControl3D(
                 value={"azimuth": 0, "elevation": 0},
-                elem_id="lighting-3d-control"
-            )
-            run_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
-            
-            gr.Markdown("### 🎚️ Slider Controls")
-            
-            azimuth_slider = gr.Slider(
-                label="Azimuth (Horizontal Rotation)",
-                minimum=0,
-                maximum=315,
-                step=45,
-                value=0,
-                info="0°=front, 90°=right, 180°=rear, 270°=left"
+                elem_id="light-3d-control"
             )
+
+            run_btn = gr.Button("✨ Render Lighting", variant="primary", size="lg")
             
-            elevation_slider = gr.Slider(
-                label="Elevation (Vertical Angle)",
-                minimum=-90,
-                maximum=90,
-                step=90,
-                value=0,
-                info="-90°=from below, 0°=horizontal, 90°=from above"
-            )
+            with gr.Accordion("🎚️ Manual Sliders", open=False):
+                azimuth_slider = gr.Slider(
+                    label="Azimuth", minimum=0, maximum=359, step=45, value=0
+                )
+                elevation_slider = gr.Slider(
+                    label="Elevation", minimum=-90, maximum=90, step=10, value=0
+                )
             
             prompt_preview = gr.Textbox(
-                label="Generated Prompt",
+                label="Active Prompt",
                 value="Light source from the Front",
                 interactive=False
             )
         
-        # Right column: Output
         with gr.Column(scale=1):
-            result = gr.Image(label="Output Image", height=500)
+            result = gr.Image(label="Relighted Result", height=600)
             
             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)
-                num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=20, step=1, value=4)
+                randomize_seed = gr.Checkbox(label="Randomize", value=True)
+                guidance_scale = gr.Slider(label="Guidance", minimum=1.0, maximum=10.0, step=0.1, value=5.0)
+                num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=20, step=1, value=4)
                 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)
+
+    # --- Wiring ---
+    def sync_3d_to_sliders(val):
+        az = val.get('azimuth', 0)
+        el = val.get('elevation', 0)
+        return az, el, build_lighting_prompt(az, el)
     
-    # --- Event Handlers ---
-    
-    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)
-            prompt = build_lighting_prompt(az, el)
-            return az, el, prompt
-        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)
-        # Convert PIL image to base64 data URL
+    def sync_sliders_to_3d(az, el):
+        return {"azimuth": az, "elevation": el}, build_lighting_prompt(az, el)
+
+    light_3d.change(sync_3d_to_sliders, [light_3d], [azimuth_slider, elevation_slider, prompt_preview])
+    azimuth_slider.change(sync_sliders_to_3d, [azimuth_slider, elevation_slider], [light_3d, prompt_preview])
+    elevation_slider.change(sync_sliders_to_3d, [azimuth_slider, elevation_slider], [light_3d, prompt_preview])
+
+    def update_on_upload(img):
+        w, h = update_dimensions_on_upload(img)
+        if img is None: 
+            return w, h, gr.update(imageUrl=None)
         import base64
         from io import BytesIO
         buffered = BytesIO()
-        image.save(buffered, format="PNG")
+        img.save(buffered, format="PNG")
         img_str = base64.b64encode(buffered.getvalue()).decode()
-        data_url = f"data:image/png;base64,{img_str}"
-        return gr.update(imageUrl=data_url)
-    
-    # Slider -> Prompt preview
-    for slider in [azimuth_slider, elevation_slider]:
-        slider.change(
-            fn=update_prompt_from_sliders,
-            inputs=[azimuth_slider, elevation_slider],
-            outputs=[prompt_preview]
-        )
-    
-    # 3D control -> Sliders + Prompt
-    lighting_3d.change(
-        fn=sync_3d_to_sliders,
-        inputs=[lighting_3d],
-        outputs=[azimuth_slider, elevation_slider, prompt_preview]
-    )
-    
-    # Sliders -> 3D control
-    for slider in [azimuth_slider, elevation_slider]:
-        slider.release(
-            fn=sync_sliders_to_3d,
-            inputs=[azimuth_slider, elevation_slider],
-            outputs=[lighting_3d]
-        )
+        return w, h, gr.update(imageUrl=f"data:image/png;base64,{img_str}")
+
+    image.upload(update_on_upload, [image], [width, height, light_3d])
     
-    # Generate button
     run_btn.click(
-        fn=infer_lighting_edit,
-        inputs=[image, azimuth_slider, elevation_slider, seed, randomize_seed, guidance_scale, num_inference_steps, height, width],
-        outputs=[result, seed, prompt_preview]
-    )
-    
-    # Image upload -> update dimensions AND update 3D preview
-    image.upload(
-        fn=update_dimensions_on_upload,
-        inputs=[image],
-        outputs=[width, height]
-    ).then(
-        fn=update_3d_image,
-        inputs=[image],
-        outputs=[lighting_3d]
+        infer_lighting_edit,
+        [image, azimuth_slider, elevation_slider, seed, randomize_seed, guidance_scale, num_inference_steps, height, width],
+        [result, seed, prompt_preview]
     )
-    
-    # Also handle image clear
-    image.clear(
-        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>'
-    css = '.fillable{max-width: 1200px !important}'
-    demo.launch(head=head, css=css)
+    # Load Three.js and OrbitControls from CDN
+    head = '''
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
+    <script src="https://cdn.jsdelivr.net/npm/three@0.128.0/examples/js/controls/OrbitControls.js"></script>
+    '''
+    demo.launch(head=head)