Update app.py

app.py

@@ -3,43 +3,42 @@ import numpy as np
 import random
 import torch
 import spaces
-import json
-import base64
-from io import BytesIO
 from PIL import Image
 
-# --- Imports based on your request ---
+# --- Updated Imports as requested ---
+from diffusers import FlowMatchEulerDiscreteScheduler
 try:
-    from diffusers import FlowMatchEulerDiscreteScheduler
-    # Attempting to import the custom pipeline/model classes
-    # Assuming these files exist in your PYTHONPATH or local directory
     from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
     from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
     from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
-except ImportError as e:
-    print(f"⚠️ Import Error: {e}")
-    print("Falling back to standard classes for UI demonstration (Generate will fail without local files).")
-    # Mocks for UI testing if files are missing
-    from diffusers import DiffusionPipeline
-    class QwenImageEditPlusPipeline(DiffusionPipeline):
-        @classmethod
-        def from_pretrained(cls, *args, **kwargs): return cls()
-        def to(self, device): return self
-        def load_lora_weights(self, *args, **kwargs): pass
-        def set_adapters(self, *args, **kwargs): pass
-        def get_active_adapters(self): return []
-        def __call__(self, *args, **kwargs): 
-            class R: images=[Image.new("RGB", (512,512), "gray")]
-            return R()
-    class QwenImageTransformer2DModel:
-        @classmethod
-        def from_pretrained(cls, *args, **kwargs): return cls()
-
-# --- Configuration ---
+except ImportError:
+    raise ImportError("Please ensure the 'qwenimage' package is installed or in your python path.")
+
 MAX_SEED = np.iinfo(np.int32).max
+
+# --- Model Configuration ---
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Load the pipeline with the custom transformer
+pipe = QwenImageEditPlusPipeline.from_pretrained(
+    "Qwen/Qwen-Image-Edit-2511",
+    transformer=QwenImageTransformer2DModel.from_pretrained(
+        "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
+        torch_dtype=dtype,
+        device_map='cuda'
+    ),
+    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}")
+
+# Adapter Specs
 ADAPTER_SPECS = {
     "Multi-Angle-Lighting": {
         "repo": "dx8152/Qwen-Edit-2509-Multi-Angle-Lighting",
@@ -48,147 +47,94 @@ ADAPTER_SPECS = {
     }
 }
 
-# --- Model Loading ---
-print("⏳ Initializing Model...")
-try:
-    transformer_model = QwenImageTransformer2DModel.from_pretrained(
-        "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
-        torch_dtype=dtype,
-        device_map=device
-    )
-
-    pipe = QwenImageEditPlusPipeline.from_pretrained(
-        "Qwen/Qwen-Image-Edit-2511",
-        transformer=transformer_model,
-        torch_dtype=dtype
-    ).to(device)
-
-    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}")
-
-except Exception as e:
-    print(f"❌ Critical Model Load Error: {e}")
-    pipe = None
+# State to track loaded adapter
+CURRENT_LOADED_ADAPTER = None
+
+# --- Prompt Building & Logic ---
+
+# Lighting mappings based on Azimuth (Horizontal)
+# 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"
+}
 
-# --- Prompt Logic ---
+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 get_lighting_prompt(azimuth: float, elevation: float) -> str:
+def build_lighting_prompt(azimuth: float, elevation: float) -> str:
     """
-    Maps Azimuth (0-360) and Elevation (-90 to 90) to the specific prompt list.
+    Build a lighting prompt based on the specific 10 options provided.
+    Priority:
+    1. If Elevation is very high (> 60) -> Above
+    2. If Elevation is very low (< -60) -> Below
+    3. Else -> Use Azimuth mapping
     """
-    # 1. Handle Vertical Extremes
-    if elevation >= 55:
+    
+    # Check Vertical extremes first
+    if elevation >= 60:
         return "Light source from Above"
-    if elevation <= -55:
+    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]
 
-    # 2. Handle Horizontal (Azimuth)
-    # Normalize azimuth to 0-360
-    az = azimuth % 360
-    
-    # 8-way split (45 degrees per sector). 
-    # Center of "Front" (0) is at 0. Sector is 337.5 to 22.5.
-    
-    if az >= 337.5 or az < 22.5:
-        return "Light source from the Front"
-    elif 22.5 <= az < 67.5:
-        return "Light source from the Right Front"
-    elif 67.5 <= az < 112.5:
-        return "Light source from the Right"
-    elif 112.5 <= az < 157.5:
-        return "Light source from the Right Rear"
-    elif 157.5 <= az < 202.5:
-        return "Light source from the Rear"
-    elif 202.5 <= az < 247.5:
-        return "Light source from the Left Rear"
-    elif 247.5 <= az < 292.5:
-        return "Light source from the Left"
-    elif 292.5 <= az < 337.5:
-        return "Light source from the Left Front"
-    
-    return "Light source from the Front" # Fallback
-
-# --- Helper Functions ---
-
-def load_adapter_if_needed():
-    """Lazy loads the LoRA only when generation is requested."""
-    if pipe is None: return
+@spaces.GPU
+def infer_lighting_edit(
+    image: Image.Image,
+    azimuth: float = 0.0,
+    elevation: float = 0.0,
+    seed: int = 0,
+    randomize_seed: bool = True,
+    guidance_scale: float = 1.0,
+    num_inference_steps: int = 4,
+    height: int = 1024,
+    width: int = 1024,
+):
+    global CURRENT_LOADED_ADAPTER
     
+    # --- Lazy Loading Logic ---
     spec = ADAPTER_SPECS["Multi-Angle-Lighting"]
-    active = pipe.get_active_adapters()
-    
-    if spec["adapter_name"] not in active:
-        print(f"📥 Lazy Loading LoRA: {spec['repo']}...")
+    if CURRENT_LOADED_ADAPTER != spec["adapter_name"]:
+        print(f"Lazy loading adapter: {spec['adapter_name']}...")
         pipe.load_lora_weights(
             spec["repo"],
             weight_name=spec["weights"],
             adapter_name=spec["adapter_name"]
         )
         pipe.set_adapters([spec["adapter_name"]], adapter_weights=[1.0])
-        print("✅ LoRA Loaded and Active.")
-
-def update_dimensions_on_upload(image):
-    if image is None: return 1024, 1024
-    w, h = image.size
-    # Aspect ratio logic to keep near 1024
-    if w > h:
-        new_w, new_h = 1024, int(1024 * (h / w))
-    else:
-        new_h, new_w = 1024, int(1024 * (w / h))
-    # Snap to multiples of 8
-    return (new_w // 8) * 8, (new_h // 8) * 8
-
-def get_image_base64(image):
-    if image is None: return None
-    buffered = BytesIO()
-    image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    return f"data:image/png;base64,{img_str}"
-
-# --- Inference ---
-
-@spaces.GPU
-def infer_lighting(
-    image: Image.Image,
-    azimuth: float,
-    elevation: float,
-    seed: int,
-    randomize_seed: bool,
-    guidance_scale: float,
-    steps: int,
-    height: int,
-    width: int,
-):
-    if pipe is None:
-        raise gr.Error("Model failed to initialize. Check console logs.")
+        CURRENT_LOADED_ADAPTER = spec["adapter_name"]
     
-    if image is None:
-        raise gr.Error("Please upload an image first.")
-
-    # 1. Ensure Adapter is Loaded
-    load_adapter_if_needed()
-
-    # 2. Build Prompt
-    prompt = get_lighting_prompt(azimuth, elevation)
-    print(f"💡 Generated Prompt: {prompt}")
+    # --- Generation ---
+    progress = gr.Progress(track_tqdm=True)
+    
+    prompt = build_lighting_prompt(azimuth, elevation)
+    print(f"Generated Lighting Prompt: {prompt}")
 
-    # 3. Setup Generator
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
-    
-    pil_image = image.convert("RGB")
 
-    # 4. Run Pipeline
+    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=f"<sks> {prompt}", # Assuming <sks> triggers the LoRA + prompt
-        height=height,
-        width=width,
-        num_inference_steps=steps,
+        prompt=prompt,
+        height=height if height != 0 else None,
+        width=width if width != 0 else None,
+        num_inference_steps=num_inference_steps,
         generator=generator,
         guidance_scale=guidance_scale,
         num_images_per_prompt=1,
@@ -196,388 +142,457 @@ def infer_lighting(
 
     return result, seed, prompt
 
-# --- 3D HTML Logic (Updated for Light Controller) ---
-THREE_JS_LOGIC = """
-<div id="light-control-wrapper" style="width: 100%; height: 500px; position: relative; background: radial-gradient(circle at center, #2a2a2a 0%, #111 100%); border-radius: 12px; overflow: hidden; border: 1px solid #333;">
-    <div id="prompt-overlay" style="position: absolute; top: 15px; left: 50%; transform: translateX(-50%); background: rgba(0,0,0,0.7); padding: 8px 16px; border-radius: 20px; font-family: sans-serif; font-size: 14px; font-weight: bold; color: #ffeb3b; border: 1px solid #ffeb3b; white-space: nowrap; z-index: 10; pointer-events: none;">
-        Initializing Light Studio...
-    </div>
-    <div style="position: absolute; bottom: 10px; right: 10px; color: #666; font-family: sans-serif; font-size: 10px; pointer-events: none;">
-        Drag the Sun • Right Click to Pan
-    </div>
-</div>
-<script>
-(function() {
-    const wrapper = document.getElementById('light-control-wrapper');
-    const promptOverlay = document.getElementById('prompt-overlay');
-    
-    window.lightControl = { updateState: null, updateTexture: null };
-
-    const initScene = () => {
-        if (typeof THREE === 'undefined') { setTimeout(initScene, 100); return; }
-
-        // 1. Scene Setup
-        const scene = new THREE.Scene();
-        // Fog to blend floor into background
-        scene.fog = new THREE.FogExp2(0x111111, 0.05);
-
-        const camera = new THREE.PerspectiveCamera(45, wrapper.clientWidth / wrapper.clientHeight, 0.1, 1000);
-        camera.position.set(0, 3, 6.5);
-        camera.lookAt(0, 0.5, 0);
-
-        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;
-        wrapper.appendChild(renderer.domElement);
-
-        // 2. Objects
-        const CENTER = new THREE.Vector3(0, 0.8, 0);
-        const ORBIT_RADIUS = 2.5;
-
-        // Grid / Floor
-        const grid = new THREE.GridHelper(10, 20, 0x444444, 0x222222);
-        scene.add(grid);
-
-        // Subject Plane (The Image)
-        function createPlaceholder() {
-            const cvs = document.createElement('canvas'); cvs.width=256; cvs.height=256;
-            const ctx = cvs.getContext('2d');
-            ctx.fillStyle = '#333'; ctx.fillRect(0,0,256,256);
-            ctx.fillStyle = '#555'; ctx.font = '30px Arial'; ctx.textAlign='center';
-            ctx.fillText("SUBJECT", 128, 128);
-            return new THREE.CanvasTexture(cvs);
-        }
+def update_dimensions_on_upload(image):
+    if image is None:
+        return 1024, 1024
+    original_width, original_height = image.size
+    if original_width > original_height:
+        new_width = 1024
+        aspect_ratio = original_height / original_width
+        new_height = int(new_width * aspect_ratio)
+    else:
+        new_height = 1024
+        aspect_ratio = original_width / original_height
+        new_width = int(new_height * aspect_ratio)
+    new_width = (new_width // 8) * 8
+    new_height = (new_height // 8) * 8
+    return new_width, new_height
+
+# --- 3D Lighting Control Component ---
+class LightControl3D(gr.HTML):
+    """
+    A 3D Lighting control component using Three.js.
+    Visualizes a Light Source (Sun/Bulb) relative to the subject.
+    """
+    def __init__(self, value=None, imageUrl=None, **kwargs):
+        if value is None:
+            value = {"azimuth": 0, "elevation": 0}
         
-        let planeMat = new THREE.MeshStandardMaterial({ 
-            map: createPlaceholder(), 
-            side: THREE.DoubleSide,
-            roughness: 0.5, metalness: 0.1 
-        });
-        let subjectPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.5, 1.5), planeMat);
-        subjectPlane.position.copy(CENTER);
-        subjectPlane.castShadow = true; 
-        subjectPlane.receiveShadow = true;
-        scene.add(subjectPlane);
-
-        // The "Sun" (Light Source Control)
-        const lightGroup = new THREE.Group();
-        scene.add(lightGroup);
-
-        // Visual Sun
-        const sunGeo = new THREE.SphereGeometry(0.2, 32, 32);
-        const sunMat = new THREE.MeshBasicMaterial({ color: 0xffeb3b });
-        const sunMesh = new THREE.Mesh(sunGeo, sunMat);
-        lightGroup.add(sunMesh);
-
-        // Glow effect (larger translucent sphere)
-        const glowMesh = new THREE.Mesh(
-            new THREE.SphereGeometry(0.35, 16, 16),
-            new THREE.MeshBasicMaterial({ color: 0xffeb3b, transparent: true, opacity: 0.3 })
-        );
-        lightGroup.add(glowMesh);
-
-        // Actual Light Source (for casting shadows in 3D view)
-        const pointLight = new THREE.PointLight(0xffffff, 1.5, 100);
-        lightGroup.add(pointLight);
-
-        // Guide Rings (To visualize sphere of influence)
-        const ringMat = new THREE.LineBasicMaterial({ color: 0x444444, transparent: true, opacity: 0.3 });
-        const eqRing = new THREE.LineLoop(new THREE.CircleGeometry(ORBIT_RADIUS, 64), ringMat);
-        eqRing.rotation.x = Math.PI/2; eqRing.position.y = CENTER.y;
-        scene.add(eqRing);
+        html_template = """
+        <div id="light-control-wrapper" style="width: 100%; height: 450px; position: relative; background: #0b0b0b; border-radius: 12px; overflow: hidden; border: 1px solid #333;">
+            <div id="prompt-overlay" style="position: absolute; top: 10px; left: 50%; transform: translateX(-50%); background: rgba(0,0,0,0.7); padding: 8px 16px; border-radius: 20px; font-family: sans-serif; font-weight: bold; font-size: 14px; color: #FFD700; white-space: nowrap; z-index: 10; border: 1px solid #FFD700;">Light Source: Front</div>
+        </div>
+        """
         
-        const merRing = new THREE.LineLoop(new THREE.CircleGeometry(ORBIT_RADIUS, 64), ringMat);
-        merRing.rotation.y = Math.PI/2; merRing.position.y = CENTER.y;
-        scene.add(merRing);
-
-        // Connector line
-        const connLine = new THREE.Line(new THREE.BufferGeometry(), new THREE.LineBasicMaterial({ color: 0xffeb3b, opacity: 0.5, transparent: true }));
-        scene.add(connLine);
-
-        // 3. Logic & State
-        let azimuth = 0;   // Degrees
-        let elevation = 0; // Degrees
-
-        const promptMap = [
-            { max: 55,  min: -999, label: "Front" }, // Placeholder logic
-        ];
-
-        function getLabel(az, el) {
-            if (el >= 55) return "LIGHT: ABOVE";
-            if (el <= -55) return "LIGHT: BELOW";
-            
-            const a = (az % 360 + 360) % 360;
-            if (a >= 337.5 || a < 22.5) return "LIGHT: FRONT";
-            if (a >= 22.5 && a < 67.5) return "LIGHT: RIGHT FRONT";
-            if (a >= 67.5 && a < 112.5) return "LIGHT: RIGHT";
-            if (a >= 112.5 && a < 157.5) return "LIGHT: RIGHT REAR";
-            if (a >= 157.5 && a < 202.5) return "LIGHT: REAR";
-            if (a >= 202.5 && a < 247.5) return "LIGHT: LEFT REAR";
-            if (a >= 247.5 && a < 292.5) return "LIGHT: LEFT";
-            if (a >= 292.5 && a < 337.5) return "LIGHT: LEFT FRONT";
-            return "LIGHT";
-        }
-
-        function updatePositions() {
-            // Convert Spherical to Cartesian
-            const phi = THREE.MathUtils.degToRad(90 - elevation); // Elevation from Y axis
-            const theta = THREE.MathUtils.degToRad(azimuth);      // Azimuth around Y axis
-
-            const x = ORBIT_RADIUS * Math.sin(phi) * Math.sin(theta);
-            const y = ORBIT_RADIUS * Math.cos(phi) + CENTER.y;
-            const z = ORBIT_RADIUS * Math.sin(phi) * Math.cos(theta);
-
-            lightGroup.position.set(x, y, z);
-            lightGroup.lookAt(CENTER);
-
-            // Update Connector
-            connLine.geometry.setFromPoints([CENTER, lightGroup.position]);
-
-            // Update Text
-            promptOverlay.innerText = getLabel(azimuth, elevation);
-            promptOverlay.style.boxShadow = `0 0 15px rgba(255, 235, 59, ${Math.max(0.2, (elevation+90)/180)})`;
-        }
-
-        // 4. Interaction (Drag the Sun)
-        const raycaster = new THREE.Raycaster();
-        const mouse = new THREE.Vector2();
-        let isDragging = false;
-        
-        // Invisible sphere for raycasting drag
-        const dragSphere = new THREE.Mesh(
-            new THREE.SphereGeometry(ORBIT_RADIUS, 32, 32),
-            new THREE.MeshBasicMaterial({ visible: false })
-        );
-        dragSphere.position.copy(CENTER);
-        scene.add(dragSphere);
-
-        function handleInput(e) {
-            if (!isDragging) return;
+        js_on_load = """
+        (() => {
+            const wrapper = element.querySelector('#light-control-wrapper');
+            const promptOverlay = element.querySelector('#prompt-overlay');
             
-            const rect = wrapper.getBoundingClientRect();
-            // Handle touch or mouse
-            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;
-
-            raycaster.setFromCamera(mouse, camera);
-            const intersects = raycaster.intersectObject(dragSphere);
-
-            if (intersects.length > 0) {
-                const point = intersects[0].point;
-                const rel = new THREE.Vector3().subVectors(point, CENTER);
+            const initScene = () => {
+                if (typeof THREE === 'undefined') {
+                    setTimeout(initScene, 100);
+                    return;
+                }
                 
-                // Cartesian to Spherical
-                const r = rel.length();
-                // Elevation (Lat)
-                const lat = 90 - THREE.MathUtils.radToDeg(Math.acos(rel.y / r));
-                // Azimuth (Lon)
-                const lon = THREE.MathUtils.radToDeg(Math.atan2(rel.x, rel.z));
-
-                elevation = Math.max(-90, Math.min(90, lat));
-                azimuth = (lon + 360) % 360; // Normalize
-
-                updatePositions();
-            }
-        }
-        
-        // Event Listeners
-        const getMouse = (e) => {
-            const rect = wrapper.getBoundingClientRect();
-            return {
-                x: ((e.clientX - rect.left) / rect.width) * 2 - 1,
-                y: -((e.clientY - rect.top) / rect.height) * 2 + 1
-            };
-        };
-
-        wrapper.addEventListener('mousedown', (e) => {
-            const m = getMouse(e);
-            raycaster.setFromCamera(m, camera);
-            const intersects = raycaster.intersectObject(sunMesh); // Hit the sun?
-            if(intersects.length > 0 || raycaster.intersectObject(glowMesh).length > 0) {
-                isDragging = true;
-                wrapper.style.cursor = 'grabbing';
-                // Disable OrbitControls usually attached to camera if we want pure sun control
-                // But for now, we assume user hits the sun specifically
-            }
-        });
-
-        window.addEventListener('mousemove', (e) => {
-            if(isDragging) {
-                handleInput(e);
-            } else {
-                // Hover state
-                const m = getMouse(e);
-                raycaster.setFromCamera(m, camera);
-                if(raycaster.intersectObject(sunMesh).length > 0) {
-                    wrapper.style.cursor = 'grab';
-                    sunMesh.scale.setScalar(1.2);
-                } else {
-                    wrapper.style.cursor = 'default';
-                    sunMesh.scale.setScalar(1.0);
+                // Scene setup
+                const scene = new THREE.Scene();
+                scene.background = new THREE.Color(0x0b0b0b);
+                scene.fog = new THREE.FogExp2(0x0b0b0b, 0.1);
+                
+                const camera = new THREE.PerspectiveCamera(50, wrapper.clientWidth / wrapper.clientHeight, 0.1, 1000);
+                camera.position.set(3.5, 2.5, 3.5);
+                camera.lookAt(0, 0.5, 0);
+                
+                const renderer = new THREE.WebGLRenderer({ antialias: true, alpha: true });
+                renderer.setSize(wrapper.clientWidth, wrapper.clientHeight);
+                renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
+                wrapper.insertBefore(renderer.domElement, promptOverlay);
+                
+                // Ambient Light (dim)
+                scene.add(new THREE.AmbientLight(0xffffff, 0.2));
+                
+                // The Grid
+                const grid = new THREE.GridHelper(6, 12, 0x444444, 0x222222);
+                scene.add(grid);
+                
+                // Constants
+                const CENTER = new THREE.Vector3(0, 0.75, 0);
+                const ORBIT_RADIUS = 2.2;
+                
+                // State
+                let azimuthAngle = props.value?.azimuth || 0;
+                let elevationAngle = props.value?.elevation || 0;
+                
+                // --- Mappings for UI ---
+                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";
+                    
+                    // Snap azimuth
+                    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];
                 }
-            }
-        });
 
-        window.addEventListener('mouseup', () => {
-            if(isDragging) {
-                isDragging = false;
-                wrapper.style.cursor = 'default';
+                // --- Objects ---
                 
-                // Snap for cleaner output? 
-                // Let's snap to nearest integer for clean sliders
-                azimuth = Math.round(azimuth);
-                elevation = Math.round(elevation);
-
-                // Send to Python
-                const bridge = document.querySelector("#bridge-output textarea");
-                if (bridge) {
-                    bridge.value = JSON.stringify({ azimuth, elevation });
-                    bridge.dispatchEvent(new Event("input", { bubbles: true }));
+                // 1. The Subject (Plane with image)
+                function createPlaceholderTexture() {
+                    const canvas = document.createElement('canvas');
+                    canvas.width = 256; canvas.height = 256;
+                    const ctx = canvas.getContext('2d');
+                    ctx.fillStyle = '#222'; ctx.fillRect(0,0,256,256);
+                    ctx.fillStyle = '#444'; ctx.beginPath(); ctx.arc(128,128, 80, 0, Math.PI*2); ctx.fill();
+                    ctx.strokeStyle = '#666'; ctx.lineWidth=4; ctx.beginPath(); ctx.moveTo(64,64); ctx.lineTo(192,192); ctx.moveTo(192,64); ctx.lineTo(64,192); ctx.stroke();
+                    return new THREE.CanvasTexture(canvas);
                 }
-            }
-        });
-
-        // 5. Render Loop
-        function animate() {
-            requestAnimationFrame(animate);
-            renderer.render(scene, camera);
-        }
-        animate();
-        updatePositions(); // Initial draw
-
-        // 6. External API
-        window.lightControl.updateState = (data) => {
-            if(typeof data === 'string') data = JSON.parse(data);
-            if(data) {
-                azimuth = data.azimuth !== undefined ? data.azimuth : azimuth;
-                elevation = data.elevation !== undefined ? data.elevation : elevation;
-                updatePositions();
-            }
-        };
+                
+                let currentTexture = createPlaceholderTexture();
+                const planeMaterial = new THREE.MeshBasicMaterial({ map: currentTexture, side: THREE.DoubleSide });
+                let targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.2, 1.2), planeMaterial);
+                targetPlane.position.copy(CENTER);
+                scene.add(targetPlane);
+                
+                function updateTextureFromUrl(url) {
+                    if (!url) {
+                        planeMaterial.map = createPlaceholderTexture();
+                        planeMaterial.needsUpdate = true;
+                        return;
+                    }
+                    new THREE.TextureLoader().load(url, (tex) => {
+                        planeMaterial.map = tex;
+                        planeMaterial.needsUpdate = true;
+                        const img = tex.image;
+                        if(img.width && img.height) {
+                            const aspect = img.width / img.height;
+                            targetPlane.scale.set(aspect > 1 ? 1 : aspect, aspect > 1 ? 1/aspect : 1, 1);
+                        }
+                    });
+                }
+                if (props.imageUrl) updateTextureFromUrl(props.imageUrl);
 
-        window.lightControl.updateTexture = (url) => {
-            if (!url) {
-                planeMat.map = createPlaceholder();
-                planeMat.needsUpdate = true;
-                return;
-            }
-            new THREE.TextureLoader().load(url, (tex) => {
-                tex.colorSpace = THREE.SRGBColorSpace;
-                planeMat.map = tex;
+                // 2. The Light Source Visualizer (Yellow Sphere)
+                const lightGroup = new THREE.Group();
+                scene.add(lightGroup);
                 
-                const img = tex.image;
-                const aspect = img.width / img.height;
-                const scale = 1.5; 
-                if (aspect > 1) subjectPlane.scale.set(scale, scale / aspect, 1);
-                else subjectPlane.scale.set(scale * aspect, scale, 1);
+                // The "Sun"
+                const sunGeo = new THREE.SphereGeometry(0.25, 32, 32);
+                const sunMat = new THREE.MeshBasicMaterial({ color: 0xFFD700 });
+                const sunMesh = new THREE.Mesh(sunGeo, sunMat);
                 
-                planeMat.needsUpdate = true;
-            });
-        };
-    };
-
-    initScene();
-})();
-</script>
-"""
-
-# --- UI Setup ---
-css = """
-#col-container { max-width: 1400px; margin: 0 auto; }
-#light-control-wrapper { box-shadow: 0 8px 32px rgba(0,0,0,0.8); border: 2px solid #222; }
-.gradio-container { background-color: #0b0f19 !important; }
-h1, p { color: #e2e8f0 !important; }
-.block.svelte-1t38q2d { background: #1a202c !important; border-color: #2d3748 !important; }
-"""
-
-with gr.Blocks() as demo:
-    gr.HTML('<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>')
-    
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown("# 💡 Qwen-Edit-2509 — Multi-Angle Lighting Studio")
-        gr.Markdown("Control the lighting direction using the **3D Sun Controller** or sliders below. The system automatically selects the correct lighting prompt based on the sun's position.")
-        
-        with gr.Row():
-            # Left: Controls
-            with gr.Column(scale=4):
-                # 3D Viewport
-                gr.HTML(THREE_JS_LOGIC)
+                // Glow effect (simple semi-transparent sphere)
+                const glowGeo = new THREE.SphereGeometry(0.4, 32, 32);
+                const glowMat = new THREE.MeshBasicMaterial({ color: 0xFFD700, transparent: true, opacity: 0.3 });
+                const glowMesh = new THREE.Mesh(glowGeo, glowMat);
+                sunMesh.add(glowMesh);
                 
-                # Hidden Bridges
-                bridge_output = gr.Textbox(elem_id="bridge-output", visible=False) # JS -> Python
-                bridge_input = gr.JSON(value={}, visible=False) # Python -> JS
-
-                with gr.Row():
-                    azimuth_slider = gr.Slider(0, 360, label="Sun Azimuth (Rotation)", value=0)
-                    elevation_slider = gr.Slider(-90, 90, label="Sun Elevation (Height)", value=0)
+                lightGroup.add(sunMesh);
+                
+                // Directional Line (Ray of light)
+                const lineGeo = new THREE.BufferGeometry().setFromPoints([new THREE.Vector3(0,0,0), new THREE.Vector3(0,0,-1)]);
+                const lineMat = new THREE.LineDashedMaterial({ color: 0xFFD700, dashSize: 0.2, gapSize: 0.1 });
+                const lightRay = new THREE.Line(lineGeo, lineMat);
+                lightRay.computeLineDistances();
+                lightGroup.add(lightRay);
+                
+                // Actual Light for the scene (to light up the 3D meshes)
+                const sceneLight = new THREE.DirectionalLight(0xffffff, 2.0);
+                lightGroup.add(sceneLight);
 
-                prompt_preview = gr.Textbox(label="Active Lighting Prompt", value="Light source from the Front", interactive=False)
+                // 3. UI Handles (Rings)
                 
-                with gr.Row():
-                    image = gr.Image(label="Subject Image", type="pil", height=200)
-                    run_btn = gr.Button("✨ Render Light", variant="primary", scale=2)
+                // Azimuth Ring (Green-ish)
+                const azRing = new THREE.Mesh(
+                    new THREE.TorusGeometry(ORBIT_RADIUS, 0.03, 16, 64),
+                    new THREE.MeshBasicMaterial({ color: 0x444444, transparent: true, opacity: 0.5 })
+                );
+                azRing.rotation.x = Math.PI/2;
+                azRing.position.y = 0.1;
+                scene.add(azRing);
+
+                // --- Update Logic ---
+                function updatePositions() {
+                    const azRad = THREE.MathUtils.degToRad(azimuthAngle);
+                    const elRad = THREE.MathUtils.degToRad(elevationAngle);
+                    
+                    // Spherical coordinates
+                    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(x, y, z);
+                    lightGroup.lookAt(CENTER);
+                    
+                    // Scale ray to reach center
+                    const dist = lightGroup.position.distanceTo(CENTER);
+                    lightRay.scale.z = dist;
+
+                    // Update UI text
+                    promptOverlay.textContent = getPromptText(azimuthAngle, elevationAngle);
+                    
+                    // Update Sun Color based on position (Warning colors for Above/Below)
+                    if (elevationAngle >= 60 || elevationAngle <= -60) {
+                        sunMat.color.setHex(0xFF4500); // Orange-Red for vertical extremes
+                        promptOverlay.style.color = "#FF4500";
+                        promptOverlay.style.borderColor = "#FF4500";
+                    } else {
+                        sunMat.color.setHex(0xFFD700); // Gold for horizontal
+                        promptOverlay.style.color = "#FFD700";
+                        promptOverlay.style.borderColor = "#FFD700";
+                    }
+                }
 
-            # Right: Result & Settings
-            with gr.Column(scale=3):
-                result = gr.Image(label="Relit Result")
+                // --- Interaction (Dragging) ---
+                const raycaster = new THREE.Raycaster();
+                const mouse = new THREE.Vector2();
+                let isDragging = false;
                 
-                with gr.Accordion("🛠️ Advanced Settings", open=False):
-                    seed = gr.Slider(0, MAX_SEED, value=42, label="Seed")
-                    randomize_seed = gr.Checkbox(True, label="Randomize Seed")
-                    guidance_scale = gr.Slider(1, 10, 5.0, step=0.1, label="Guidance Scale")
-                    steps = gr.Slider(1, 50, 20, step=1, label="Inference Steps")
-                    width = gr.Slider(256, 2048, 1024, step=8, label="Width")
-                    height = gr.Slider(256, 2048, 1024, step=8, label="Height")
+                // Invisible sphere for raycasting drag
+                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);
+
+                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.intersectObject(sunMesh); // Click strictly on sun
+                    if(intersects.length > 0) {
+                        isDragging = true;
+                        wrapper.style.cursor = 'grabbing';
+                        // Disable orbit controls if we added them, but here we are manual
+                    }
+                }
 
-    # --- Event Wiring ---
+                function onMove(e) {
+                    if(!isDragging) {
+                         // Hover effect
+                        const rect = wrapper.getBoundingClientRect();
+                        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.intersectObject(sunMesh);
+                        wrapper.style.cursor = intersects.length > 0 ? 'grab' : 'default';
+                        return;
+                    }
+                    
+                    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);
+                    const intersects = raycaster.intersectObject(dragSphere);
+                    
+                    if(intersects.length > 0) {
+                        const point = intersects[0].point;
+                        const rel = new THREE.Vector3().subVectors(point, CENTER);
+                        
+                        // Calculate Azimuth
+                        let az = Math.atan2(rel.x, rel.z) * (180/Math.PI);
+                        if(az < 0) az += 360;
+                        
+                        // Calculate Elevation
+                        const dist = Math.sqrt(rel.x*rel.x + rel.z*rel.z);
+                        let el = Math.atan2(rel.y, dist) * (180/Math.PI);
+                        
+                        // Clamp Elevation slightly to avoid gimbal lock flip visual issues
+                        el = Math.max(-89, Math.min(89, el));
+                        
+                        azimuthAngle = az;
+                        elevationAngle = el;
+                        updatePositions();
+                    }
+                }
 
-    # 1. Update Prompt Display
-    def update_prompt_ui(az, el):
-        return get_lighting_prompt(az, el)
+                function onUp() {
+                    if(isDragging) {
+                        isDragging = false;
+                        wrapper.style.cursor = 'default';
+                        
+                        // Snap logic for output
+                        let finalAz = azimuthAngle;
+                        let finalEl = elevationAngle;
+                        
+                        // Simple snap visual helper could go here, but we just trigger values
+                        props.value = { azimuth: finalAz, elevation: finalEl };
+                        trigger('change', props.value);
+                    }
+                }
 
-    # 2. Image Handling
-    def handle_image_upload(img):
-        w, h = update_dimensions_on_upload(img)
-        b64 = get_image_base64(img)
-        return w, h, b64
+                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);
 
-    image.upload(handle_image_upload, inputs=[image], outputs=[width, height, bridge_input]) \
-         .then(None, [image], None, js="(img) => { if(img) window.lightControl.updateTexture(img); }")
+                // Initial render
+                updatePositions();
+                
+                function animate() {
+                    requestAnimationFrame(animate);
+                    renderer.render(scene, camera);
+                }
+                animate();
+                
+                // Watchers
+                setInterval(() => {
+                    if (props.imageUrl !== currentTexture.sourceFile) { // simple check
+                         // handled by prop check mostly
+                    }
+                    // Sync from Sliders
+                    if (props.value && (props.value.azimuth !== azimuthAngle || props.value.elevation !== elevationAngle)) {
+                        if(!isDragging) {
+                            azimuthAngle = props.value.azimuth;
+                            elevationAngle = props.value.elevation;
+                            updatePositions();
+                        }
+                    }
+                }, 100);
+                
+                wrapper._updateTexture = updateTextureFromUrl;
+            }
+            initScene();
+        })();
+        """
+        
+        super().__init__(
+            value=value,
+            html_template=html_template,
+            js_on_load=js_on_load,
+            imageUrl=imageUrl,
+            **kwargs
+        )
 
-    # 3. Sliders -> Bridge -> 3D View
-    def sync_sliders_to_bridge(az, el):
-        return {"azimuth": az, "elevation": el}
+# --- UI Layout ---
+css = '''
+#col-container { max-width: 1200px; margin: 0 auto; }
+.dark .progress-text { color: white !important; }
+#light-3d-control { min-height: 450px; }
+'''
 
-    azimuth_slider.change(sync_sliders_to_bridge, [azimuth_slider, elevation_slider], bridge_input)
-    elevation_slider.change(sync_sliders_to_bridge, [azimuth_slider, elevation_slider], bridge_input)
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 💡 Qwen Edit 2509 — Multi-Angle Lighting Control
     
-    # Trigger JS update when bridge input changes
-    bridge_input.change(None, [bridge_input], None, js="(val) => window.lightControl.updateState(val)")
+    Control the **Light Source Direction** relative to your subject.
+    Drag the **Yellow Sun** in the 3D viewport or use the sliders to position the light.
+    """)
+    
+    with gr.Row():
+        # Left column: Input and Controls
+        with gr.Column(scale=1):
+            image = gr.Image(label="Input Image", type="pil", height=300)
+            
+            gr.Markdown("### 🎮 3D Light Source Controller")
+            
+            light_3d = LightControl3D(
+                value={"azimuth": 0, "elevation": 0},
+                elem_id="light-3d-control"
+            )
 
-    # Update Prompt when sliders move
-    azimuth_slider.change(update_prompt_ui, [azimuth_slider, elevation_slider], prompt_preview)
-    elevation_slider.change(update_prompt_ui, [azimuth_slider, elevation_slider], prompt_preview)
+            run_btn = gr.Button("🚀 Generate Lighting", variant="primary", size="lg")
+            
+            gr.Markdown("### 🎚️ Fine-Tune Position")
+            
+            azimuth_slider = gr.Slider(
+                label="Light Azimuth (Horizontal)", 
+                minimum=0, maximum=359, step=45, value=0,
+                info="0°=Front, 90°=Right, 180°=Rear, 270°=Left"
+            )
+            
+            elevation_slider = gr.Slider(
+                label="Light Elevation (Vertical)", 
+                minimum=-90, maximum=90, step=10, value=0,
+                info=">60° = Above, <-60° = Below"
+            )
+            
+            prompt_preview = gr.Textbox(
+                label="Active Lighting Prompt",
+                value="Light source from the Front",
+                interactive=False
+            )
+        
+        # Right column: Output
+        with gr.Column(scale=1):
+            result = gr.Image(label="Relighted Result", height=500)
+            
+            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=5.0)
+                num_inference_steps = gr.Slider(label="Inference 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)
 
-    # 4. 3D View (Bridge Output) -> Sliders
-    def sync_bridge_to_sliders(data_str):
-        try:
-            data = json.loads(data_str)
-            return data.get('azimuth', 0), data.get('elevation', 0)
-        except:
-            return 0, 0
+    # --- Event Wiring ---
 
-    bridge_output.change(sync_bridge_to_sliders, bridge_output, [azimuth_slider, elevation_slider])
+    def sync_3d_to_sliders(val):
+        """Syncs 3D controller movement to sliders and prompt"""
+        az = val.get('azimuth', 0)
+        el = val.get('elevation', 0)
+        prompt = build_lighting_prompt(az, el)
+        return az, el, prompt
+    
+    def sync_sliders_to_3d(az, el):
+        """Syncs slider movement to 3D controller and prompt"""
+        prompt = build_lighting_prompt(az, el)
+        return {"azimuth": az, "elevation": el}, prompt
+
+    # 3D -> Sliders
+    light_3d.change(
+        fn=sync_3d_to_sliders,
+        inputs=[light_3d],
+        outputs=[azimuth_slider, elevation_slider, prompt_preview]
+    )
+    
+    # Sliders -> 3D
+    azimuth_slider.change(
+        fn=sync_sliders_to_3d,
+        inputs=[azimuth_slider, elevation_slider],
+        outputs=[light_3d, prompt_preview]
+    )
+    elevation_slider.change(
+        fn=sync_sliders_to_3d,
+        inputs=[azimuth_slider, elevation_slider],
+        outputs=[light_3d, prompt_preview]
+    )
 
-    # 5. Generation
+    # Upload Image handling
+    def update_on_upload(img):
+        w, h = update_dimensions_on_upload(img)
+        # Convert to data URL for 3D texture
+        if img is None: 
+            return w, h, gr.update(imageUrl=None)
+        
+        import base64
+        from io import BytesIO
+        buffered = BytesIO()
+        img.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+        data_url = f"data:image/png;base64,{img_str}"
+        return w, h, gr.update(imageUrl=data_url)
+
+    image.upload(
+        fn=update_on_upload,
+        inputs=[image],
+        outputs=[width, height, light_3d]
+    )
+    
+    # Generate
     run_btn.click(
-        infer_lighting,
-        inputs=[image, azimuth_slider, elevation_slider, seed, randomize_seed, guidance_scale, steps, height, width],
+        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]
     )
 
 if __name__ == "__main__":
-    demo.launch(css=css, theme=gr.themes.Base())
+    head = '<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>'
+    demo.launch(head=head, css=css)