Update app.py

app.py

@@ -9,9 +9,11 @@ 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
+
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
+# --- Theme Configuration ---
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -80,10 +82,11 @@ class OrangeRedTheme(Soft):
 
 orange_red_theme = OrangeRedTheme()
 
+# --- Global Constants & Model Setup ---
 MAX_SEED = np.iinfo(np.int32).max
+
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2511",
     transformer=QwenImageTransformer2DModel.from_pretrained(
@@ -93,7 +96,6 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
     ),
     torch_dtype=dtype
 ).to(device)
-
 try:
     pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
     print("Flash Attention 3 Processor set successfully.")
@@ -107,7 +109,6 @@ ADAPTER_SPECS = {
         "adapter_name": "multi-angle-lighting"
     },
 }
-
 loaded = False
 
 AZIMUTH_MAP = {
@@ -127,12 +128,19 @@ ELEVATION_MAP = {
     90: "Above"
 }
 
+# --- Helper Functions ---
+
 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.
+    """
     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:
@@ -151,6 +159,8 @@ def infer_lighting_edit(
     width: int = 1024,
 ):
     global loaded
+    progress = gr.Progress(track_tqdm=True)
+    
     if not loaded:
         pipe.load_lora_weights(
             ADAPTER_SPECS["Multi-Angle-Lighting"]["repo"],
@@ -159,20 +169,15 @@ def infer_lighting_edit(
         )
         pipe.set_adapters([ADAPTER_SPECS["Multi-Angle-Lighting"]["adapter_name"]], adapter_weights=[1.0])
         loaded = True
-
+    
     prompt = build_lighting_prompt(azimuth, elevation)
     print(f"Generated Prompt: {prompt}")
-
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-
     generator = torch.Generator(device=device).manual_seed(seed)
-
     if image is None:
         raise gr.Error("Please upload an image first.")
-
     pil_image = image.convert("RGB") if isinstance(image, Image.Image) else Image.open(image).convert("RGB")
-
     result = pipe(
         image=[pil_image],
         prompt=prompt,
@@ -183,7 +188,6 @@ 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):
@@ -202,61 +206,76 @@ def update_dimensions_on_upload(image):
     new_height = (new_height // 8) * 8
     return new_width, new_height
 
+# --- Custom 3D Component ---
+
 class LightingControl3D(gr.HTML):
+    """
+    A 3D lighting control component using Three.js.
+    """
     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>
         """
-
+        
         js_on_load = """
         (() => {
             const wrapper = element.querySelector('#lighting-control-wrapper');
             const promptOverlay = element.querySelector('#prompt-overlay');
-
+            
+            // Wait for THREE to load
             const initScene = () => {
                 if (typeof THREE === 'undefined') {
                     setTimeout(initScene, 100);
                     return;
                 }
-
+                
+                // Scene setup
                 const scene = new THREE.Scene();
                 scene.background = new THREE.Color(0x1a1a1a);
-
+                
                 const camera = new THREE.PerspectiveCamera(50, wrapper.clientWidth / wrapper.clientHeight, 0.1, 1000);
-                camera.position.set(5, 3.5, 5);
-                camera.lookAt(0, 0.8, 0);
-
+                camera.position.set(4.5, 3, 4.5);
+                camera.lookAt(0, 0.75, 0);
+                
                 const renderer = new THREE.WebGLRenderer({ antialias: true });
                 renderer.setSize(wrapper.clientWidth, wrapper.clientHeight);
                 renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
                 renderer.shadowMap.enabled = true;
                 renderer.shadowMap.type = THREE.PCFSoftShadowMap;
                 wrapper.insertBefore(renderer.domElement, promptOverlay);
-
-                scene.add(new THREE.AmbientLight(0xffffff, 0.15));
-
+                
+                // Lighting (Ambient)
+                scene.add(new THREE.AmbientLight(0xffffff, 0.1));
+                
+                // Ground plane for shadows
                 const ground = new THREE.Mesh(
-                    new THREE.PlaneGeometry(12, 12),
-                    new THREE.ShadowMaterial({ opacity: 0.35 })
+                    new THREE.PlaneGeometry(10, 10),
+                    new THREE.ShadowMaterial({ opacity: 0.3 })
                 );
                 ground.rotation.x = -Math.PI / 2;
-                ground.position.y = -0.01;
+                ground.position.y = 0;
                 ground.receiveShadow = true;
                 scene.add(ground);
-
-                scene.add(new THREE.GridHelper(10, 20, 0x333333, 0x222222));
-
-                const CENTER = new THREE.Vector3(0, 0.8, 0);
-                const BASE_DISTANCE = 3.2;
-
+                
+                // Grid
+                scene.add(new THREE.GridHelper(8, 16, 0x333333, 0x222222));
+                
+                // Constants
+                const CENTER = new THREE.Vector3(0, 0.75, 0);
+                const BASE_DISTANCE = 2.5;
+                const AZIMUTH_RADIUS = 2.4;
+                const ELEVATION_RADIUS = 1.8;
+                
+                // State
                 let azimuthAngle = props.value?.azimuth || 0;
                 let elevationAngle = props.value?.elevation || 0;
-
+                
+                // Mappings
                 const azimuthSteps = [0, 45, 90, 135, 180, 225, 270, 315];
                 const elevationSteps = [-90, 0, 90];
                 const azimuthNames = {
@@ -265,11 +284,12 @@ class LightingControl3D(gr.HTML):
                     270: 'Left', 315: 'Left Front'
                 };
                 const elevationNames = { '-90': 'Below', '0': '', '90': 'Above' };
-
+                
                 function snapToNearest(value, steps) {
                     return steps.reduce((prev, curr) => Math.abs(curr - value) < Math.abs(prev - value) ? curr : prev);
                 }
-
+                
+                // Create placeholder texture (smiley face)
                 function createPlaceholderTexture() {
                     const canvas = document.createElement('canvas');
                     canvas.width = 256;
@@ -293,25 +313,28 @@ class LightingControl3D(gr.HTML):
                     ctx.stroke();
                     return new THREE.CanvasTexture(canvas);
                 }
-
+                
+                // Target image plane
                 let currentTexture = createPlaceholderTexture();
-                const planeMaterial = new THREE.MeshStandardMaterial({
-                    map: currentTexture,
-                    side: THREE.DoubleSide,
-                    roughness: 0.6,
-                    metalness: 0.05
-                });
-                let targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.5, 1.5), planeMaterial);
+                const planeMaterial = new THREE.MeshStandardMaterial({ map: currentTexture, side: THREE.DoubleSide, roughness: 0.5, metalness: 0 });
+                let targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.2, 1.2), planeMaterial);
                 targetPlane.position.copy(CENTER);
                 targetPlane.receiveShadow = true;
                 scene.add(targetPlane);
-
+                
+                // Function to update texture from image URL
                 function updateTextureFromUrl(url) {
                     if (!url) {
                         planeMaterial.map = createPlaceholderTexture();
                         planeMaterial.needsUpdate = true;
+                        scene.remove(targetPlane);
+                        targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.2, 1.2), planeMaterial);
+                        targetPlane.position.copy(CENTER);
+                        targetPlane.receiveShadow = true;
+                        scene.add(targetPlane);
                         return;
                     }
+                    
                     const loader = new THREE.TextureLoader();
                     loader.crossOrigin = 'anonymous';
                     loader.load(url, (texture) => {
@@ -319,14 +342,16 @@ class LightingControl3D(gr.HTML):
                         texture.magFilter = THREE.LinearFilter;
                         planeMaterial.map = texture;
                         planeMaterial.needsUpdate = true;
-
+                        
                         const img = texture.image;
                         if (img && img.width && img.height) {
                             const aspect = img.width / img.height;
-                            const maxSize = 1.8;
-                            let planeWidth = maxSize;
-                            let planeHeight = maxSize / aspect;
-                            if (aspect < 1) {
+                            const maxSize = 1.5;
+                            let planeWidth, planeHeight;
+                            if (aspect > 1) {
+                                planeWidth = maxSize;
+                                planeHeight = maxSize / aspect;
+                            } else {
                                 planeHeight = maxSize;
                                 planeWidth = maxSize * aspect;
                             }
@@ -339,205 +364,208 @@ class LightingControl3D(gr.HTML):
                             targetPlane.receiveShadow = true;
                             scene.add(targetPlane);
                         }
+                    }, undefined, (err) => {
+                        console.error('Failed to load texture:', err);
                     });
                 }
-
+                
                 if (props.imageUrl) {
                     updateTextureFromUrl(props.imageUrl);
                 }
-
-                // ── Redesigned Softbox (red housing, white soft light) ──
-                const softboxGroup = new THREE.Group();
-
-                // Main softbox body (red)
-                const softboxGeo = new THREE.BoxGeometry(1.8, 1.2, 0.25);
-                const softboxMat = new THREE.MeshStandardMaterial({
-                    color: 0xc41e3a,  // red
-                    roughness: 0.7,
-                    metalness: 0.3
-                });
-                const softboxBody = new THREE.Mesh(softboxGeo, softboxMat);
-                softboxGroup.add(softboxBody);
-
-                // Diffusion panel (white glowing)
-                const diffusionGeo = new THREE.PlaneGeometry(1.68, 1.08);
-                const diffusionMat = new THREE.MeshStandardMaterial({
-                    color: 0xffffff,
-                    emissive: 0xffffff,
-                    emissiveIntensity: 4.5,
-                    roughness: 0.9,
-                    metalness: 0.0,
-                    side: THREE.FrontSide
-                });
-                const diffusion = new THREE.Mesh(diffusionGeo, diffusionMat);
-                diffusion.position.z = 0.126;
-                softboxGroup.add(diffusion);
-
-                // Simple barn doors (red)
-                const doorMat = softboxMat;
-                const topDoor = new THREE.Mesh(new THREE.BoxGeometry(1.8, 0.4, 0.04), doorMat);
-                topDoor.position.set(0, 0.6, 0.13);
-                topDoor.rotation.x = THREE.MathUtils.degToRad(110);
-                softboxGroup.add(topDoor);
-
-                const bottomDoor = new THREE.Mesh(new THREE.BoxGeometry(1.8, 0.4, 0.04), doorMat);
-                bottomDoor.position.set(0, -0.6, 0.13);
-                bottomDoor.rotation.x = THREE.MathUtils.degToRad(-110);
-                softboxGroup.add(bottomDoor);
-
-                softboxGroup.position.set(0, 0, 0);
-                scene.add(softboxGroup);
-
-                // SpotLight for stronger cast shadow
-                const spotLight = new THREE.SpotLight(0xffffff, 12, 12, Math.PI / 3.5, 0.4, 1.5);
-                spotLight.position.set(0, 0, 0.15);
+                
+                // --- LIGHT MODEL REDESIGN (Softbox) ---
+                const lightGroup = new THREE.Group();
+                
+                // Box Geometry for Softbox
+                const sbGeometry = new THREE.BoxGeometry(0.6, 0.6, 0.3);
+                
+                // Materials: Red body, White light face
+                const sbMatRed = new THREE.MeshStandardMaterial({ color: 0xff0000, roughness: 0.5 });
+                const sbMatWhite = new THREE.MeshStandardMaterial({ color: 0xffffff, emissive: 0xffffff, emissiveIntensity: 2 });
+                
+                // In Three.js, Face 4 is +Z. When we lookAt(), +Z points to target.
+                // 0:+x, 1:-x, 2:+y, 3:-y, 4:+z, 5:-z
+                const sbMaterials = [
+                    sbMatRed, // Right
+                    sbMatRed, // Left
+                    sbMatRed, // Top
+                    sbMatRed, // Bottom
+                    sbMatWhite, // Front (Light Face)
+                    sbMatRed  // Back
+                ];
+                
+                const softbox = new THREE.Mesh(sbGeometry, sbMaterials);
+                lightGroup.add(softbox);
+                
+                // Actual SpotLight source
+                const spotLight = new THREE.SpotLight(0xffffff, 10, 10, Math.PI / 3, 1, 1);
+                spotLight.position.set(0, 0, 0.1); // Slightly in front of the box center
                 spotLight.castShadow = true;
-                spotLight.shadow.mapSize.width = 2048;
-                spotLight.shadow.mapSize.height = 2048;
-                spotLight.shadow.camera.near = 0.1;
-                spotLight.shadow.camera.far = 50;
-                spotLight.shadow.bias = -0.0005;
-                softboxGroup.add(spotLight);
-
+                spotLight.shadow.mapSize.width = 1024;
+                spotLight.shadow.mapSize.height = 1024;
+                spotLight.shadow.camera.near = 0.5;
+                spotLight.shadow.camera.far = 500;
+                spotLight.shadow.bias = -0.005;
+                lightGroup.add(spotLight);
+                
                 const lightTarget = new THREE.Object3D();
                 lightTarget.position.copy(CENTER);
                 scene.add(lightTarget);
                 spotLight.target = lightTarget;
-
-                // ── Controls ──
+                
+                scene.add(lightGroup);
+                
+                // --- HANDLES REDESIGN ---
+                
+                // YELLOW: Azimuth ring
                 const azimuthRing = new THREE.Mesh(
-                    new THREE.TorusGeometry(2.8, 0.05, 16, 64),
-                    new THREE.MeshStandardMaterial({ color: 0x444444, emissive: 0x222222 })
+                    new THREE.TorusGeometry(AZIMUTH_RADIUS, 0.04, 16, 64),
+                    new THREE.MeshStandardMaterial({ color: 0xffff00, emissive: 0xffff00, emissiveIntensity: 0.3 })
                 );
                 azimuthRing.rotation.x = Math.PI / 2;
                 azimuthRing.position.y = 0.05;
                 scene.add(azimuthRing);
-
+                
                 const azimuthHandle = new THREE.Mesh(
-                    new THREE.SphereGeometry(0.22, 24, 24),
-                    new THREE.MeshStandardMaterial({ color: 0xffcc00, emissive: 0xffcc00, emissiveIntensity: 0.7 })  // Yellow
+                    new THREE.SphereGeometry(0.18, 16, 16),
+                    new THREE.MeshStandardMaterial({ color: 0xffff00, emissive: 0xffff00, emissiveIntensity: 0.5 })
                 );
                 azimuthHandle.userData.type = 'azimuth';
                 scene.add(azimuthHandle);
-
+                
+                // BLUE: 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)));
+                }
+                const arcCurve = new THREE.CatmullRomCurve3(arcPoints);
                 const elevationArc = new THREE.Mesh(
-                    new THREE.TorusGeometry(1.8, 0.05, 16, 64, Math.PI),
-                    new THREE.MeshStandardMaterial({ color: 0x007bff, emissive: 0x007bff, emissiveIntensity: 0.4 })  // Blue arc
+                    new THREE.TubeGeometry(arcCurve, 32, 0.04, 8, false),
+                    new THREE.MeshStandardMaterial({ color: 0x0000ff, emissive: 0x0000ff, emissiveIntensity: 0.3 })
                 );
-                elevationArc.rotation.z = Math.PI / 2;
-                elevationArc.position.set(-1.0, CENTER.y, 0);
                 scene.add(elevationArc);
-
+                
                 const elevationHandle = new THREE.Mesh(
-                    new THREE.SphereGeometry(0.22, 24, 24),
-                    new THREE.MeshStandardMaterial({ color: 0x007bff, emissive: 0x007bff, emissiveIntensity: 0.7 })  // Blue
+                    new THREE.SphereGeometry(0.18, 16, 16),
+                    new THREE.MeshStandardMaterial({ color: 0x0000ff, emissive: 0x0000ff, emissiveIntensity: 0.5 })
                 );
                 elevationHandle.userData.type = 'elevation';
                 scene.add(elevationHandle);
-
-                // Reset Button inside 3D view
-                const resetBtn = document.createElement('div');
-                resetBtn.innerHTML = 'Reset';
-                resetBtn.style.position = 'absolute';
-                resetBtn.style.top = '12px';
-                resetBtn.style.right = '12px';
-                resetBtn.style.background = 'rgba(220, 53, 69, 0.85)';
-                resetBtn.style.color = 'white';
-                resetBtn.style.padding = '8px 14px';
-                resetBtn.style.borderRadius = '8px';
-                resetBtn.style.cursor = 'pointer';
-                resetBtn.style.zIndex = '15';
-                resetBtn.style.fontSize = '14px';
-                resetBtn.style.fontWeight = 'bold';
-                resetBtn.style.boxShadow = '0 2px 8px rgba(0,0,0,0.5)';
-                wrapper.appendChild(resetBtn);
-
-                resetBtn.addEventListener('click', () => {
+                
+                // --- RESET BUTTON REDESIGN ---
+                const refreshBtn = document.createElement('button');
+                refreshBtn.innerHTML = 'Reset View';
+                refreshBtn.style.position = 'absolute';
+                refreshBtn.style.top = '15px';
+                refreshBtn.style.right = '15px';
+                refreshBtn.style.background = 'rgba(255, 255, 255, 0.9)';
+                refreshBtn.style.border = 'none';
+                refreshBtn.style.borderRadius = '20px'; // Pill shape
+                refreshBtn.style.padding = '8px 16px';
+                refreshBtn.style.fontFamily = 'sans-serif';
+                refreshBtn.style.fontWeight = 'bold';
+                refreshBtn.style.fontSize = '12px';
+                refreshBtn.style.color = '#333';
+                refreshBtn.style.cursor = 'pointer';
+                refreshBtn.style.zIndex = '20';
+                refreshBtn.style.boxShadow = '0 2px 5px rgba(0,0,0,0.3)';
+                refreshBtn.style.transition = 'background 0.2s';
+                
+                refreshBtn.onmouseover = () => { refreshBtn.style.background = '#ffffff'; };
+                refreshBtn.onmouseout = () => { refreshBtn.style.background = 'rgba(255, 255, 255, 0.9)'; };
+                
+                wrapper.appendChild(refreshBtn);
+                
+                refreshBtn.addEventListener('click', () => {
                     azimuthAngle = 0;
                     elevationAngle = 0;
                     updatePositions();
                     updatePropsAndTrigger();
                 });
-
+                
                 function updatePositions() {
+                    const distance = BASE_DISTANCE;
                     const azRad = THREE.MathUtils.degToRad(azimuthAngle);
                     const elRad = THREE.MathUtils.degToRad(elevationAngle);
-
-                    const lightX = BASE_DISTANCE * Math.sin(azRad) * Math.cos(elRad);
-                    const lightY = BASE_DISTANCE * Math.sin(elRad) + CENTER.y;
-                    const lightZ = BASE_DISTANCE * Math.cos(azRad) * Math.cos(elRad);
-
-                    softboxGroup.position.set(lightX, lightY, lightZ);
-                    softboxGroup.lookAt(CENTER);
-
-                    azimuthHandle.position.set(2.8 * Math.sin(azRad), 0.05, 2.8 * Math.cos(azRad));
-                    elevationHandle.position.set(-1.0, 1.8 * Math.sin(elRad) + CENTER.y, 1.8 * Math.cos(elRad));
-
+                    
+                    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);
+                    
+                    lightGroup.position.set(lightX, lightY, lightZ);
+                    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 promptText = 'Light source from';
+                    let prompt = 'Light source from';
                     if (elSnap !== 0) {
-                        promptText += ' ' + elevationNames[String(elSnap)];
+                        prompt += ' ' + elevationNames[String(elSnap)];
                     } else {
-                        promptText += ' the ' + azimuthNames[azSnap];
+                        prompt += ' the ' + azimuthNames[azSnap];
                     }
-                    promptOverlay.textContent = promptText;
+                    promptOverlay.textContent = prompt;
                 }
-
+                
                 function updatePropsAndTrigger() {
                     const azSnap = snapToNearest(azimuthAngle, azimuthSteps);
                     const elSnap = snapToNearest(elevationAngle, elevationSteps);
+                    
                     props.value = { azimuth: azSnap, elevation: elSnap };
                     trigger('change', props.value);
                 }
-
+                
+                // Raycasting
                 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;
-
+                
                 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;
-
+                    
                     raycaster.setFromCamera(mouse, camera);
                     const intersects = raycaster.intersectObjects([azimuthHandle, elevationHandle]);
-
+                    
                     if (intersects.length > 0) {
                         isDragging = true;
                         dragTarget = intersects[0].object;
-                        dragTarget.material.emissiveIntensity = 1.2;
-                        dragTarget.scale.setScalar(1.4);
+                        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;
-
+                    
                     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, new THREE.Vector3())) {
-                                const intersect = new THREE.Vector3();
-                                raycaster.ray.intersectPlane(plane, intersect);
-                                azimuthAngle = THREE.MathUtils.radToDeg(Math.atan2(intersect.x, intersect.z));
+                            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), -1.0);
-                            if (raycaster.ray.intersectPlane(plane, new THREE.Vector3())) {
-                                const intersect = new THREE.Vector3();
-                                raycaster.ray.intersectPlane(plane, intersect);
-                                const relY = intersect.y - CENTER.y;
-                                const relZ = intersect.z;
+                            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);
                             }
                         }
@@ -546,43 +574,44 @@ class LightingControl3D(gr.HTML):
                         raycaster.setFromCamera(mouse, camera);
                         const intersects = raycaster.intersectObjects([azimuthHandle, elevationHandle]);
                         [azimuthHandle, elevationHandle].forEach(h => {
-                            h.material.emissiveIntensity = 0.7;
+                            h.material.emissiveIntensity = 0.5;
                             h.scale.setScalar(1);
                         });
                         if (intersects.length > 0) {
-                            intersects[0].object.material.emissiveIntensity = 1.0;
-                            intersects[0].object.scale.setScalar(1.2);
+                            intersects[0].object.material.emissiveIntensity = 0.8;
+                            intersects[0].object.scale.setScalar(1.1);
                             canvas.style.cursor = 'grab';
                         } else {
                             canvas.style.cursor = 'default';
                         }
                     }
                 });
-
+                
                 const onMouseUp = () => {
                     if (dragTarget) {
-                        dragTarget.material.emissiveIntensity = 0.7;
+                        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);
-
+                            
                             let azDiff = targetAz - startAz;
                             if (azDiff > 180) azDiff -= 360;
                             if (azDiff < -180) azDiff += 360;
                             azimuthAngle = startAz + azDiff * ease;
                             if (azimuthAngle < 0) azimuthAngle += 360;
                             if (azimuthAngle >= 360) azimuthAngle -= 360;
-
+                            
                             elevationAngle = startEl + (targetEl - startEl) * ease;
-
+                            
                             updatePositions();
                             if (t < 1) requestAnimationFrame(animateSnap);
                             else updatePropsAndTrigger();
@@ -593,75 +622,83 @@ class LightingControl3D(gr.HTML):
                     dragTarget = null;
                     canvas.style.cursor = 'default';
                 };
-
+                
                 canvas.addEventListener('mouseup', onMouseUp);
                 canvas.addEventListener('mouseleave', onMouseUp);
-
-                // Touch events (simplified)
+                // Touch support
                 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.2;
-                        dragTarget.scale.setScalar(1.4);
+                        dragTarget.material.emissiveIntensity = 1.0;
+                        dragTarget.scale.setScalar(1.3);
+                        dragStartMouse.copy(mouse);
                     }
                 }, { passive: false });
-
+                
                 canvas.addEventListener('touchmove', (e) => {
                     e.preventDefault();
                     const touch = e.touches[0];
                     const rect = canvas.getBoundingClientRect();
                     mouse.x = ((touch.clientX - rect.left) / rect.width) * 2 - 1;
                     mouse.y = -((touch.clientY - rect.top) / rect.height) * 2 + 1;
-
+                    
                     if (isDragging && dragTarget) {
                         raycaster.setFromCamera(mouse, camera);
-                        // same logic as mousemove...
                         if (dragTarget.userData.type === 'azimuth') {
                             const plane = new THREE.Plane(new THREE.Vector3(0, 1, 0), -0.05);
-                            const intersect = new THREE.Vector3();
-                            if (raycaster.ray.intersectPlane(plane, intersect)) {
-                                azimuthAngle = THREE.MathUtils.radToDeg(Math.atan2(intersect.x, intersect.z));
+                            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), -1.0);
-                            const intersect = new THREE.Vector3();
-                            if (raycaster.ray.intersectPlane(plane, intersect)) {
-                                const relY = intersect.y - CENTER.y;
-                                const relZ = intersect.z;
+                            const plane = new THREE.Plane(new THREE.Vector3(1, 0, 0), -0.8);
+                            if (raycaster.ray.intersectPlane(plane, intersection)) {
+                                const relY = intersection.y - CENTER.y;
+                                const relZ = intersection.z;
                                 elevationAngle = THREE.MathUtils.clamp(THREE.MathUtils.radToDeg(Math.atan2(relY, relZ)), -90, 90);
                             }
                         }
                         updatePositions();
                     }
                 }, { passive: false });
-
-                canvas.addEventListener('touchend', onMouseUp, { passive: false });
-
+                
+                canvas.addEventListener('touchend', (e) => {
+                    e.preventDefault();
+                    onMouseUp();
+                }, { passive: false });
+                canvas.addEventListener('touchcancel', (e) => {
+                    e.preventDefault();
+                    onMouseUp();
+                }, { passive: false });
+                
+                // Initial update
                 updatePositions();
-
+                
+                // Render loop
                 function render() {
                     requestAnimationFrame(render);
                     renderer.render(scene, camera);
                 }
                 render();
-
+                
+                // Handle resize
                 new ResizeObserver(() => {
                     camera.aspect = wrapper.clientWidth / wrapper.clientHeight;
                     camera.updateProjectionMatrix();
                     renderer.setSize(wrapper.clientWidth, wrapper.clientHeight);
                 }).observe(wrapper);
-
+                
+                // Store update functions for external calls
                 wrapper._updateFromProps = (newVal) => {
                     if (newVal && typeof newVal === 'object') {
                         azimuthAngle = newVal.azimuth ?? azimuthAngle;
@@ -669,9 +706,10 @@ class LightingControl3D(gr.HTML):
                         updatePositions();
                     }
                 };
-
+                
                 wrapper._updateTexture = updateTextureFromUrl;
-
+                
+                // Watch for prop changes
                 let lastImageUrl = props.imageUrl;
                 let lastValue = JSON.stringify(props.value);
                 setInterval(() => {
@@ -690,11 +728,11 @@ class LightingControl3D(gr.HTML):
                     }
                 }, 100);
             };
-
+            
             initScene();
         })();
         """
-
+        
         super().__init__(
             value=value,
             html_template=html_template,
@@ -703,6 +741,8 @@ class LightingControl3D(gr.HTML):
             **kwargs
         )
 
+# --- Gradio UI Layout ---
+
 css = '''
 #col-container { max-width: 1200px; margin: 0 auto; }
 .dark .progress-text { color: white !important; }
@@ -710,25 +750,25 @@ css = '''
 .slider-row { display: flex; gap: 10px; align-items: center; }
 #main-title h1 {font-size: 2.4em !important;}
 '''
-
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# **Qwen-Image-Edit-2511-3D-Lighting-Control**", elem_id="main-title")
     gr.Markdown("Control lighting directions using the **3D viewport** or **sliders**. Use [Multi-Angle-Lighting](https://huggingface.co/dx8152/Qwen-Edit-2509-Multi-Angle-Lighting) LoRA for precise lighting control, paired with [Rapid-AIO-V19](https://huggingface.co/prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19).")
+    
     with gr.Row():
         with gr.Column(scale=1):
             image = gr.Image(label="Input Image", type="pil", height=300)
-
+            
             gr.Markdown("### 3D Lighting Control")
-            gr.Markdown("*Drag the handles: 🟡 Azimuth (yellow), 🔵 Elevation (blue)*")
-
+            gr.Markdown("*Drag the colored handles: 🟡 Azimuth (Direction), 🔵 Elevation (Height)*")
+            
             lighting_3d = LightingControl3D(
                 value={"azimuth": 0, "elevation": 0},
                 elem_id="lighting-3d-control"
             )
             run_btn = gr.Button("Generate Image", variant="primary", size="lg")
-
+            
             gr.Markdown("### Slider Controls")
-
+            
             azimuth_slider = gr.Slider(
                 label="Azimuth (Horizontal Rotation)",
                 minimum=0,
@@ -737,7 +777,7 @@ with gr.Blocks(css=css) as demo:
                 value=0,
                 info="0°=front, 90°=right, 180°=rear, 270°=left"
             )
-
+            
             elevation_slider = gr.Slider(
                 label="Elevation (Vertical Angle)",
                 minimum=-90,
@@ -746,6 +786,7 @@ with gr.Blocks(css=css) as demo:
                 value=0,
                 info="-90°=from below, 0°=horizontal, 90°=from above"
             )
+
             with gr.Row():
                 prompt_preview = gr.Textbox(
                     label="Generated Prompt",
@@ -753,10 +794,10 @@ with gr.Blocks(css=css) as demo:
                     interactive=True,
                     lines=1,
                 )
-
+        
         with gr.Column(scale=1):
             result = gr.Image(label="Output Image", 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)
@@ -764,23 +805,27 @@ with gr.Blocks(css=css) as demo:
                 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)
-
+    
     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)
         import base64
@@ -790,33 +835,33 @@ with gr.Blocks(css=css) as demo:
         img_str = base64.b64encode(buffered.getvalue()).decode()
         data_url = f"data:image/png;base64,{img_str}"
         return gr.update(imageUrl=data_url)
-
+    
     for slider in [azimuth_slider, elevation_slider]:
         slider.change(
             fn=update_prompt_from_sliders,
             inputs=[azimuth_slider, elevation_slider],
             outputs=[prompt_preview]
         )
-
+    
     lighting_3d.change(
         fn=sync_3d_to_sliders,
         inputs=[lighting_3d],
         outputs=[azimuth_slider, elevation_slider, prompt_preview]
     )
-
+    
     for slider in [azimuth_slider, elevation_slider]:
         slider.release(
             fn=sync_sliders_to_3d,
             inputs=[azimuth_slider, elevation_slider],
             outputs=[lighting_3d]
         )
-
+    
     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(
         fn=update_dimensions_on_upload,
         inputs=[image],
@@ -826,12 +871,12 @@ with gr.Blocks(css=css) as demo:
         inputs=[image],
         outputs=[lighting_3d]
     )
-
+    
     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}'