Update app.py

app.py

@@ -13,9 +13,6 @@ from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
-# -----------------------------
-# THEME
-# -----------------------------
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -31,7 +28,6 @@ colors.orange_red = colors.Color(
     c950="#802200",
 )
 
-
 class OrangeRedTheme(Soft):
     def __init__(
         self,
@@ -41,14 +37,10 @@ class OrangeRedTheme(Soft):
         neutral_hue: colors.Color | str = colors.slate,
         text_size: sizes.Size | str = sizes.text_lg,
         font: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("Outfit"),
-            "Arial",
-            "sans-serif",
+            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
         ),
         font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("IBM Plex Mono"),
-            "ui-monospace",
-            "monospace",
+            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
         ),
     ):
         super().__init__(
@@ -87,27 +79,21 @@ class OrangeRedTheme(Soft):
             block_label_background_fill="*primary_200",
         )
 
-
 orange_red_theme = OrangeRedTheme()
 
-# -----------------------------
-# 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(
         "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
         torch_dtype=dtype,
-        device_map="cuda",
+        device_map='cuda'
     ),
-    torch_dtype=dtype,
+    torch_dtype=dtype
 ).to(device)
-
 try:
     pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
     print("Flash Attention 3 Processor set successfully.")
@@ -117,8 +103,8 @@ except Exception as e:
 ADAPTER_SPECS = {
     "Multi-Angle-Lighting": {
         "repo": "dx8152/Qwen-Edit-2509-Multi-Angle-Lighting",
-        "weights": "多角度灯光-251121.safetensors",
-        "adapter_name": "multi-angle-lighting",
+        "weights": "多角度灯光-251116.safetensors",
+        "adapter_name": "multi-angle-lighting"
     },
 }
 loaded = False
@@ -131,30 +117,39 @@ AZIMUTH_MAP = {
     180: "Rear",
     225: "Left Rear",
     270: "Left",
-    315: "Left Front",
+    315: "Left Front"
 }
 
 ELEVATION_MAP = {
     -90: "Below",
     0: "",
-    90: "Above",
+    90: "Above"
 }
 
-
 def snap_to_nearest(value, options):
+    """Snap a value to the nearest option in a list."""
     return min(options, key=lambda x: abs(x - value))
 
-
 def build_lighting_prompt(azimuth: float, elevation: float) -> str:
+    """
+    Build a lighting prompt from azimuth and elevation values.
+    
+    Args:
+        azimuth: Horizontal rotation in degrees (0-360)
+        elevation: Vertical angle in degrees (-90 to 90)
+    
+    Returns:
+        Formatted prompt string for the LoRA
+    """
+    # Snap to nearest valid values
     azimuth_snapped = snap_to_nearest(azimuth, list(AZIMUTH_MAP.keys()))
     elevation_snapped = snap_to_nearest(elevation, list(ELEVATION_MAP.keys()))
-
+    
     if elevation_snapped == 0:
         return f"Light source from the {AZIMUTH_MAP[azimuth_snapped]}"
     else:
         return f"Light source from {ELEVATION_MAP[elevation_snapped]}"
 
-
 @spaces.GPU
 def infer_lighting_edit(
     image: Image.Image,
@@ -167,37 +162,29 @@ 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)
-
+    
     if not loaded:
         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"],
-        )
-        pipe.set_adapters(
-            [ADAPTER_SPECS["Multi-Angle-Lighting"]["adapter_name"]],
-            adapter_weights=[1.0],
+            adapter_name=ADAPTER_SPECS["Multi-Angle-Lighting"]["adapter_name"]
         )
+        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")
-    )
-
+    pil_image = image.convert("RGB") if isinstance(image, Image.Image) else Image.open(image).convert("RGB")
     result = pipe(
         image=[pil_image],
         prompt=prompt,
@@ -208,11 +195,10 @@ def infer_lighting_edit(
         guidance_scale=guidance_scale,
         num_images_per_prompt=1,
     ).images[0]
-
     return result, seed, prompt
 
-
 def update_dimensions_on_upload(image):
+    """Compute recommended dimensions preserving aspect ratio."""
     if image is None:
         return 1024, 1024
     original_width, original_height = image.size
@@ -228,390 +214,589 @@ def update_dimensions_on_upload(image):
     new_height = (new_height // 8) * 8
     return new_width, new_height
 
-
-# -----------------------------
-# CUSTOM 3D LIGHTING CONTROL
-# -----------------------------
 class LightingControl3D(gr.HTML):
     """
-    3D lighting control using Three.js
-    - Softbox: red body, white light
-    - Azimuth handle: yellow
-    - Elevation handle: blue
+    A 3D lighting control component using Three.js.
+    Outputs: { azimuth: number, elevation: number }
+    Accepts imageUrl prop to display user's uploaded image on the plane.
     """
-
     def __init__(self, value=None, imageUrl=None, **kwargs):
         if value is None:
             value = {"azimuth": 0, "elevation": 0}
-
+        
         html_template = """
-        <div id="lighting-control-wrapper" style="width: 100%; height: 450px; position: relative; background: #1a1a1a; border-radius: 12px; overflow: hidden;">
-            <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="lighting-control-wrapper" style="width: 100%; height: 450px; position: relative; background: linear-gradient(180deg, #0d1117 0%, #161b22 50%, #21262d 100%); border-radius: 12px; overflow: hidden; border: 1px solid #30363d;">
+            <div id="prompt-overlay" style="position: absolute; bottom: 10px; left: 50%; transform: translateX(-50%); background: rgba(0,0,0,0.85); padding: 10px 20px; border-radius: 8px; font-family: 'Courier New', monospace; font-size: 13px; color: #58a6ff; white-space: nowrap; z-index: 10; border: 1px solid #30363d; box-shadow: 0 4px 12px rgba(0,0,0,0.4);"></div>
+            <div id="controls-legend" style="position: absolute; top: 10px; left: 10px; background: rgba(0,0,0,0.75); padding: 8px 12px; border-radius: 6px; font-family: sans-serif; font-size: 11px; color: #c9d1d9; z-index: 10; border: 1px solid #30363d;">
+                <div style="display: flex; align-items: center; gap: 6px; margin-bottom: 4px;">
+                    <span style="width: 12px; height: 12px; background: #ffd700; border-radius: 50%; display: inline-block;"></span>
+                    <span>Azimuth (Direction)</span>
+                </div>
+                <div style="display: flex; align-items: center; gap: 6px;">
+                    <span style="width: 12px; height: 12px; background: #4dabf7; border-radius: 50%; display: inline-block;"></span>
+                    <span>Elevation (Height)</span>
+                </div>
+            </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
-                );
+                scene.background = new THREE.Color(0x0d1117);
+                
+                // Add fog for depth
+                scene.fog = new THREE.Fog(0x0d1117, 8, 20);
+                
+                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 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);
-
-                // Lighting
-                scene.add(new THREE.AmbientLight(0xffffff, 0.15));
-
-                // Ground plane for shadows
-                const ground = new THREE.Mesh(
-                    new THREE.PlaneGeometry(10, 10),
-                    new THREE.ShadowMaterial({ opacity: 0.3 })
-                );
+                renderer.toneMapping = THREE.ACESFilmicToneMapping;
+                renderer.toneMappingExposure = 1.2;
+                wrapper.insertBefore(renderer.domElement, wrapper.firstChild);
+                
+                // Ambient lighting
+                const ambientLight = new THREE.AmbientLight(0x404050, 0.3);
+                scene.add(ambientLight);
+                
+                // Ground plane with better material
+                const groundGeometry = new THREE.PlaneGeometry(12, 12);
+                const groundMaterial = new THREE.MeshStandardMaterial({ 
+                    color: 0x1a1a2e,
+                    roughness: 0.8,
+                    metalness: 0.2
+                });
+                const ground = new THREE.Mesh(groundGeometry, groundMaterial);
                 ground.rotation.x = -Math.PI / 2;
                 ground.position.y = 0;
                 ground.receiveShadow = true;
                 scene.add(ground);
-
-                // Grid
-                scene.add(new THREE.GridHelper(8, 16, 0x333333, 0x222222));
-
+                
+                // Shadow-only plane
+                const shadowPlane = new THREE.Mesh(
+                    new THREE.PlaneGeometry(12, 12),
+                    new THREE.ShadowMaterial({ opacity: 0.4 })
+                );
+                shadowPlane.rotation.x = -Math.PI / 2;
+                shadowPlane.position.y = 0.001;
+                shadowPlane.receiveShadow = true;
+                scene.add(shadowPlane);
+                
+                // Grid with better styling
+                const gridHelper = new THREE.GridHelper(10, 20, 0x30363d, 0x21262d);
+                gridHelper.position.y = 0.002;
+                scene.add(gridHelper);
+                
                 // 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 = {
-                    0: 'Front',
-                    45: 'Right Front',
-                    90: 'Right',
-                    135: 'Right Rear',
-                    180: 'Rear',
-                    225: 'Left Rear',
-                    270: 'Left',
-                    315: 'Left Front'
+                    0: 'Front', 45: 'Right Front', 90: 'Right',
+                    135: 'Right Rear', 180: 'Rear', 225: 'Left Rear',
+                    270: 'Left', 315: 'Left Front'
                 };
                 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
-                    );
+                    return steps.reduce((prev, curr) => Math.abs(curr - value) < Math.abs(prev - value) ? curr : prev);
                 }
-
-                // Placeholder texture for target image plane
+                
+                // Create placeholder texture (camera icon style)
                 function createPlaceholderTexture() {
                     const canvas = document.createElement('canvas');
                     canvas.width = 256;
                     canvas.height = 256;
                     const ctx = canvas.getContext('2d');
-                    ctx.fillStyle = '#3a3a4a';
+                    
+                    // Background gradient
+                    const gradient = ctx.createLinearGradient(0, 0, 256, 256);
+                    gradient.addColorStop(0, '#2d333b');
+                    gradient.addColorStop(1, '#22272e');
+                    ctx.fillStyle = gradient;
                     ctx.fillRect(0, 0, 256, 256);
-                    ctx.fillStyle = '#ffcc99';
+                    
+                    // Border
+                    ctx.strokeStyle = '#444c56';
+                    ctx.lineWidth = 4;
+                    ctx.strokeRect(10, 10, 236, 236);
+                    
+                    // Camera icon
+                    ctx.fillStyle = '#58a6ff';
+                    ctx.beginPath();
+                    ctx.roundRect(78, 98, 100, 70, 8);
+                    ctx.fill();
+                    
+                    // Lens
+                    ctx.fillStyle = '#2d333b';
                     ctx.beginPath();
-                    ctx.arc(128, 128, 80, 0, Math.PI * 2);
+                    ctx.arc(128, 133, 25, 0, Math.PI * 2);
                     ctx.fill();
-                    ctx.fillStyle = '#333';
+                    
+                    ctx.fillStyle = '#58a6ff';
                     ctx.beginPath();
-                    ctx.arc(100, 110, 10, 0, Math.PI * 2);
-                    ctx.arc(156, 110, 10, 0, Math.PI * 2);
+                    ctx.arc(128, 133, 18, 0, Math.PI * 2);
                     ctx.fill();
-                    ctx.strokeStyle = '#333';
-                    ctx.lineWidth = 3;
+                    
+                    // Flash
+                    ctx.fillStyle = '#58a6ff';
                     ctx.beginPath();
-                    ctx.arc(128, 130, 35, 0.2, Math.PI - 0.2);
-                    ctx.stroke();
+                    ctx.roundRect(138, 88, 30, 15, 3);
+                    ctx.fill();
+                    
+                    // Text
+                    ctx.fillStyle = '#8b949e';
+                    ctx.font = '16px sans-serif';
+                    ctx.textAlign = 'center';
+                    ctx.fillText('Upload Image', 128, 200);
+                    
                     return new THREE.CanvasTexture(canvas);
                 }
-
+                
                 // Target image plane
                 let currentTexture = createPlaceholderTexture();
-                const planeMaterial = new THREE.MeshStandardMaterial({
-                    map: currentTexture,
-                    side: THREE.DoubleSide,
-                    roughness: 0.5,
-                    metalness: 0
+                const planeMaterial = new THREE.MeshStandardMaterial({ 
+                    map: currentTexture, 
+                    side: THREE.DoubleSide, 
+                    roughness: 0.4, 
+                    metalness: 0.1 
                 });
-                let targetPlane = new THREE.Mesh(
-                    new THREE.PlaneGeometry(1.2, 1.2),
-                    planeMaterial
-                );
+                let targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.2, 1.2), planeMaterial);
                 targetPlane.position.copy(CENTER);
                 targetPlane.receiveShadow = true;
+                targetPlane.castShadow = true;
                 scene.add(targetPlane);
-
-                // Update texture from image URL
+                
+                // Frame around the image
+                const frameGeometry = new THREE.BoxGeometry(1.3, 1.3, 0.05);
+                const frameMaterial = new THREE.MeshStandardMaterial({ 
+                    color: 0x30363d, 
+                    roughness: 0.5, 
+                    metalness: 0.3 
+                });
+                const frame = new THREE.Mesh(frameGeometry, frameMaterial);
+                frame.position.set(CENTER.x, CENTER.y, -0.03);
+                scene.add(frame);
+                
+                // Function to update texture from image URL
                 function updateTextureFromUrl(url) {
                     if (!url) {
-                        // Reset to placeholder
                         planeMaterial.map = createPlaceholderTexture();
                         planeMaterial.needsUpdate = true;
-
                         scene.remove(targetPlane);
-                        targetPlane = new THREE.Mesh(
-                            new THREE.PlaneGeometry(1.2, 1.2),
-                            planeMaterial
-                        );
+                        scene.remove(frame);
+                        targetPlane = new THREE.Mesh(new THREE.PlaneGeometry(1.2, 1.2), planeMaterial);
                         targetPlane.position.copy(CENTER);
                         targetPlane.receiveShadow = true;
+                        targetPlane.castShadow = true;
                         scene.add(targetPlane);
+                        
+                        frame.geometry = new THREE.BoxGeometry(1.3, 1.3, 0.05);
+                        frame.position.set(CENTER.x, CENTER.y, -0.03);
+                        scene.add(frame);
                         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;
-
-                            const img = texture.image;
-                            if (img && 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
-                                );
-                                targetPlane.position.copy(CENTER);
-                                targetPlane.receiveShadow = true;
-                                scene.add(targetPlane);
+                    loader.load(url, (texture) => {
+                        texture.minFilter = THREE.LinearFilter;
+                        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.5;
+                            let planeWidth, planeHeight;
+                            if (aspect > 1) {
+                                planeWidth = maxSize;
+                                planeHeight = maxSize / aspect;
+                            } else {
+                                planeHeight = maxSize;
+                                planeWidth = maxSize * aspect;
                             }
-                        },
-                        undefined,
-                        (err) => {
-                            console.error('Failed to load texture:', err);
+                            scene.remove(targetPlane);
+                            scene.remove(frame);
+                            
+                            targetPlane = new THREE.Mesh(
+                                new THREE.PlaneGeometry(planeWidth, planeHeight),
+                                planeMaterial
+                            );
+                            targetPlane.position.copy(CENTER);
+                            targetPlane.receiveShadow = true;
+                            targetPlane.castShadow = true;
+                            scene.add(targetPlane);
+                            
+                            frame.geometry = new THREE.BoxGeometry(planeWidth + 0.1, planeHeight + 0.1, 0.05);
+                            frame.position.set(CENTER.x, CENTER.y, -0.03);
+                            scene.add(frame);
                         }
-                    );
+                    }, undefined, (err) => {
+                        console.error('Failed to load texture:', err);
+                    });
                 }
-
-                // Initial image, if provided
+                
                 if (props.imageUrl) {
                     updateTextureFromUrl(props.imageUrl);
                 }
-
-                // --------------------------
-                // RED SOFTBOX STUDIO LIGHT
-                // --------------------------
+                
+                // =============================================
+                // SOFTBOX LIGHT - Red box with white diffuser
+                // =============================================
                 const lightGroup = new THREE.Group();
-
-                // Softbox body (red housing)
-                const bodyGeom = new THREE.BoxGeometry(0.9, 0.7, 0.2);
-                const bodyMat = new THREE.MeshStandardMaterial({
-                    color: 0xcc2233,
-                    roughness: 0.6,
-                    metalness: 0.2
+                
+                // Main softbox frame (red)
+                const softboxDepth = 0.4;
+                const softboxWidth = 0.7;
+                const softboxHeight = 0.7;
+                
+                // Back panel (red)
+                const backPanel = new THREE.Mesh(
+                    new THREE.BoxGeometry(softboxWidth, softboxHeight, 0.05),
+                    new THREE.MeshStandardMaterial({ 
+                        color: 0xcc2222, 
+                        roughness: 0.7, 
+                        metalness: 0.3 
+                    })
+                );
+                backPanel.position.z = softboxDepth / 2;
+                lightGroup.add(backPanel);
+                
+                // Side panels (red, tapered)
+                const sideGeometry = new THREE.BoxGeometry(0.04, softboxHeight, softboxDepth);
+                const sideMaterial = new THREE.MeshStandardMaterial({ 
+                    color: 0xaa1111, 
+                    roughness: 0.6, 
+                    metalness: 0.2 
                 });
-                const bodyMesh = new THREE.Mesh(bodyGeom, bodyMat);
-                lightGroup.add(bodyMesh);
-
-                // Front panel (white light)
-                const panelGeom = new THREE.PlaneGeometry(0.7, 0.5);
-                const panelMat = new THREE.MeshStandardMaterial({
+                
+                const leftSide = new THREE.Mesh(sideGeometry, sideMaterial);
+                leftSide.position.set(-softboxWidth/2 + 0.02, 0, 0);
+                lightGroup.add(leftSide);
+                
+                const rightSide = new THREE.Mesh(sideGeometry, sideMaterial);
+                rightSide.position.set(softboxWidth/2 - 0.02, 0, 0);
+                lightGroup.add(rightSide);
+                
+                // Top and bottom panels (red)
+                const tbGeometry = new THREE.BoxGeometry(softboxWidth, 0.04, softboxDepth);
+                
+                const topPanel = new THREE.Mesh(tbGeometry, sideMaterial);
+                topPanel.position.set(0, softboxHeight/2 - 0.02, 0);
+                lightGroup.add(topPanel);
+                
+                const bottomPanel = new THREE.Mesh(tbGeometry, sideMaterial);
+                bottomPanel.position.set(0, -softboxHeight/2 + 0.02, 0);
+                lightGroup.add(bottomPanel);
+                
+                // White diffuser front (emissive white)
+                const diffuserMaterial = new THREE.MeshStandardMaterial({
                     color: 0xffffff,
                     emissive: 0xffffff,
-                    emissiveIntensity: 3.0,
-                    side: THREE.FrontSide,
-                    roughness: 0.3,
-                    metalness: 0.0
-                });
-                const panelMesh = new THREE.Mesh(panelGeom, panelMat);
-                panelMesh.position.z = 0.11;
-                lightGroup.add(panelMesh);
-
-                // Simple yoke/arm under the softbox
-                const armGeom = new THREE.CylinderGeometry(0.03, 0.03, 0.4, 16);
-                const armMat = new THREE.MeshStandardMaterial({
-                    color: 0x555555,
-                    roughness: 0.5,
-                    metalness: 0.5
+                    emissiveIntensity: 2.0,
+                    roughness: 0.2,
+                    metalness: 0,
+                    transparent: true,
+                    opacity: 0.95
                 });
-                const armMesh = new THREE.Mesh(armGeom, armMat);
-                armMesh.position.set(0, -0.6, 0);
-                lightGroup.add(armMesh);
-
-                // Spotlight (actual light) - white
-                const spotLight = new THREE.SpotLight(
-                    0xffffff,
-                    10,
-                    10,
-                    Math.PI / 3,
-                    1,
-                    1
+                
+                const diffuser = new THREE.Mesh(
+                    new THREE.PlaneGeometry(softboxWidth - 0.08, softboxHeight - 0.08),
+                    diffuserMaterial
                 );
-                // Place light slightly in front of panel
-                spotLight.position.set(0, 0, 0.15);
+                diffuser.position.z = -softboxDepth / 2 + 0.02;
+                lightGroup.add(diffuser);
+                
+                // Inner glow effect
+                const innerGlow = new THREE.Mesh(
+                    new THREE.PlaneGeometry(softboxWidth - 0.15, softboxHeight - 0.15),
+                    new THREE.MeshBasicMaterial({
+                        color: 0xffffff,
+                        transparent: true,
+                        opacity: 0.8
+                    })
+                );
+                innerGlow.position.z = -softboxDepth / 2 + 0.01;
+                lightGroup.add(innerGlow);
+                
+                // Mounting bracket
+                const bracket = new THREE.Mesh(
+                    new THREE.CylinderGeometry(0.03, 0.03, 0.15, 8),
+                    new THREE.MeshStandardMaterial({ 
+                        color: 0x333333, 
+                        roughness: 0.5, 
+                        metalness: 0.8 
+                    })
+                );
+                bracket.rotation.x = Math.PI / 2;
+                bracket.position.z = softboxDepth / 2 + 0.075;
+                lightGroup.add(bracket);
+                
+                // Ball joint
+                const ballJoint = new THREE.Mesh(
+                    new THREE.SphereGeometry(0.05, 16, 16),
+                    new THREE.MeshStandardMaterial({ 
+                        color: 0x222222, 
+                        roughness: 0.4, 
+                        metalness: 0.9 
+                    })
+                );
+                ballJoint.position.z = softboxDepth / 2 + 0.15;
+                lightGroup.add(ballJoint);
+                
+                // SpotLight for actual lighting
+                const spotLight = new THREE.SpotLight(0xffffff, 15, 12, Math.PI / 4, 0.5, 1);
+                spotLight.position.set(0, 0, -softboxDepth / 2);
                 spotLight.castShadow = true;
                 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;
+                spotLight.shadow.camera.far = 15;
+                spotLight.shadow.bias = -0.002;
                 lightGroup.add(spotLight);
-
+                
                 const lightTarget = new THREE.Object3D();
                 lightTarget.position.copy(CENTER);
                 scene.add(lightTarget);
                 spotLight.target = lightTarget;
-
+                
+                // Point light for softer fill
+                const pointLight = new THREE.PointLight(0xffffff, 3, 8);
+                pointLight.position.set(0, 0, -softboxDepth / 2 - 0.1);
+                lightGroup.add(pointLight);
+                
                 scene.add(lightGroup);
-
-                // --------------------------
-                // CONTROLLERS
-                // --------------------------
-
-                // AZIMUTH RING & HANDLE (YELLOW)
+                
+                // =============================================
+                // YELLOW: Azimuth ring and handle
+                // =============================================
                 const azimuthRing = new THREE.Mesh(
-                    new THREE.TorusGeometry(AZIMUTH_RADIUS, 0.04, 16, 64),
-                    new THREE.MeshStandardMaterial({
-                        color: 0xffd54f,          // yellow
-                        emissive: 0xffd54f,
-                        emissiveIntensity: 0.25
+                    new THREE.TorusGeometry(AZIMUTH_RADIUS, 0.035, 16, 64),
+                    new THREE.MeshStandardMaterial({ 
+                        color: 0xffd700, 
+                        emissive: 0xffd700, 
+                        emissiveIntensity: 0.4,
+                        roughness: 0.3,
+                        metalness: 0.6
                     })
                 );
                 azimuthRing.rotation.x = Math.PI / 2;
                 azimuthRing.position.y = 0.05;
                 scene.add(azimuthRing);
-
+                
+                // Azimuth direction markers
+                azimuthSteps.forEach(angle => {
+                    const rad = THREE.MathUtils.degToRad(angle);
+                    const marker = new THREE.Mesh(
+                        new THREE.SphereGeometry(0.06, 8, 8),
+                        new THREE.MeshStandardMaterial({ 
+                            color: 0xffd700, 
+                            emissive: 0xffd700, 
+                            emissiveIntensity: 0.2 
+                        })
+                    );
+                    marker.position.set(
+                        AZIMUTH_RADIUS * Math.sin(rad),
+                        0.05,
+                        AZIMUTH_RADIUS * Math.cos(rad)
+                    );
+                    scene.add(marker);
+                });
+                
                 const azimuthHandle = new THREE.Mesh(
-                    new THREE.SphereGeometry(0.18, 16, 16),
-                    new THREE.MeshStandardMaterial({
-                        color: 0xffd54f,
-                        emissive: 0xffd54f,
-                        emissiveIntensity: 0.6
+                    new THREE.SphereGeometry(0.16, 24, 24),
+                    new THREE.MeshStandardMaterial({ 
+                        color: 0xffd700, 
+                        emissive: 0xffd700, 
+                        emissiveIntensity: 0.6,
+                        roughness: 0.2,
+                        metalness: 0.7
                     })
                 );
                 azimuthHandle.userData.type = 'azimuth';
                 scene.add(azimuthHandle);
-
-                // ELEVATION ARC & HANDLE (BLUE)
+                
+                // =============================================
+                // BLUE: Elevation arc and handle
+                // =============================================
                 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)
-                        )
-                    );
+                for (let i = 0; i <= 48; i++) {
+                    const angle = THREE.MathUtils.degToRad(-90 + (180 * i / 48));
+                    arcPoints.push(new THREE.Vector3(
+                        -0.8, 
+                        ELEVATION_RADIUS * Math.sin(angle) + CENTER.y, 
+                        ELEVATION_RADIUS * Math.cos(angle)
+                    ));
                 }
                 const arcCurve = new THREE.CatmullRomCurve3(arcPoints);
                 const elevationArc = new THREE.Mesh(
-                    new THREE.TubeGeometry(arcCurve, 32, 0.04, 8, false),
-                    new THREE.MeshStandardMaterial({
-                        color: 0x2196f3,       // blue
-                        emissive: 0x2196f3,
-                        emissiveIntensity: 0.25
+                    new THREE.TubeGeometry(arcCurve, 48, 0.035, 12, false),
+                    new THREE.MeshStandardMaterial({ 
+                        color: 0x4dabf7, 
+                        emissive: 0x4dabf7, 
+                        emissiveIntensity: 0.4,
+                        roughness: 0.3,
+                        metalness: 0.6
                     })
                 );
                 scene.add(elevationArc);
-
+                
+                // Elevation markers
+                elevationSteps.forEach(angle => {
+                    const rad = THREE.MathUtils.degToRad(angle);
+                    const marker = new THREE.Mesh(
+                        new THREE.SphereGeometry(0.06, 8, 8),
+                        new THREE.MeshStandardMaterial({ 
+                            color: 0x4dabf7, 
+                            emissive: 0x4dabf7, 
+                            emissiveIntensity: 0.2 
+                        })
+                    );
+                    marker.position.set(
+                        -0.8,
+                        ELEVATION_RADIUS * Math.sin(rad) + CENTER.y,
+                        ELEVATION_RADIUS * Math.cos(rad)
+                    );
+                    scene.add(marker);
+                });
+                
                 const elevationHandle = new THREE.Mesh(
-                    new THREE.SphereGeometry(0.18, 16, 16),
-                    new THREE.MeshStandardMaterial({
-                        color: 0x2196f3,
-                        emissive: 0x2196f3,
-                        emissiveIntensity: 0.6
+                    new THREE.SphereGeometry(0.16, 24, 24),
+                    new THREE.MeshStandardMaterial({ 
+                        color: 0x4dabf7, 
+                        emissive: 0x4dabf7, 
+                        emissiveIntensity: 0.6,
+                        roughness: 0.2,
+                        metalness: 0.7
                     })
                 );
                 elevationHandle.userData.type = 'elevation';
                 scene.add(elevationHandle);
-
-                // RESET BUTTON (inside 3D control)
+                
+                // =============================================
+                // REFRESH BUTTON (styled button, not symbol)
+                // =============================================
                 const refreshBtn = document.createElement('button');
-                refreshBtn.textContent = 'Reset Light';
-                refreshBtn.style.position = 'absolute';
-                refreshBtn.style.top = '12px';
-                refreshBtn.style.right = '12px';
-                refreshBtn.style.background = 'rgba(0,0,0,0.7)';
-                refreshBtn.style.color = '#ffffff';
-                refreshBtn.style.padding = '6px 12px';
-                refreshBtn.style.borderRadius = '6px';
-                refreshBtn.style.cursor = 'pointer';
-                refreshBtn.style.zIndex = '10';
-                refreshBtn.style.fontSize = '12px';
-                refreshBtn.style.border = '1px solid rgba(255,255,255,0.2)';
-                refreshBtn.style.fontFamily = 'system-ui, -apple-system, BlinkMacSystemFont, sans-serif';
+                refreshBtn.innerHTML = '<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2.5" stroke-linecap="round" stroke-linejoin="round"><path d="M3 12a9 9 0 0 1 9-9 9.75 9.75 0 0 1 6.74 2.74L21 8"/><path d="M21 3v5h-5"/><path d="M21 12a9 9 0 0 1-9 9 9.75 9.75 0 0 1-6.74-2.74L3 16"/><path d="M3 21v-5h5"/></svg><span style="margin-left: 6px;">Reset</span>';
+                refreshBtn.style.cssText = `
+                    position: absolute;
+                    top: 10px;
+                    right: 10px;
+                    background: linear-gradient(135deg, #238636 0%, #2ea043 100%);
+                    color: white;
+                    border: none;
+                    padding: 8px 14px;
+                    border-radius: 6px;
+                    cursor: pointer;
+                    z-index: 10;
+                    font-size: 13px;
+                    font-weight: 500;
+                    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
+                    display: flex;
+                    align-items: center;
+                    box-shadow: 0 2px 8px rgba(0,0,0,0.3);
+                    transition: all 0.2s ease;
+                `;
+                refreshBtn.onmouseenter = () => {
+                    refreshBtn.style.background = 'linear-gradient(135deg, #2ea043 0%, #3fb950 100%)';
+                    refreshBtn.style.transform = 'translateY(-1px)';
+                    refreshBtn.style.boxShadow = '0 4px 12px rgba(0,0,0,0.4)';
+                };
+                refreshBtn.onmouseleave = () => {
+                    refreshBtn.style.background = 'linear-gradient(135deg, #238636 0%, #2ea043 100%)';
+                    refreshBtn.style.transform = 'translateY(0)';
+                    refreshBtn.style.boxShadow = '0 2px 8px rgba(0,0,0,0.3)';
+                };
                 wrapper.appendChild(refreshBtn);
-
+                
                 refreshBtn.addEventListener('click', () => {
-                    azimuthAngle = 0;
-                    elevationAngle = 0;
-                    updatePositions();
-                    updatePropsAndTrigger();
+                    // Animate reset
+                    const startAz = azimuthAngle;
+                    const startEl = elevationAngle;
+                    const startTime = Date.now();
+                    
+                    function animateReset() {
+                        const t = Math.min((Date.now() - startTime) / 400, 1);
+                        const ease = 1 - Math.pow(1 - t, 4);
+                        
+                        let azDiff = 0 - 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 + (0 - startEl) * ease;
+                        
+                        updatePositions();
+                        
+                        if (t < 1) {
+                            requestAnimationFrame(animateReset);
+                        } else {
+                            azimuthAngle = 0;
+                            elevationAngle = 0;
+                            updatePositions();
+                            updatePropsAndTrigger();
+                        }
+                    }
+                    animateReset();
                 });
-
-                // --------------------------
-                // POSITION / PROMPT UPDATE
-                // --------------------------
+                
                 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);
-
+                    
                     lightGroup.position.set(lightX, lightY, lightZ);
                     lightGroup.lookAt(CENTER);
-
+                    
                     azimuthHandle.position.set(
-                        AZIMUTH_RADIUS * Math.sin(azRad),
-                        0.05,
+                        AZIMUTH_RADIUS * Math.sin(azRad), 
+                        0.05, 
                         AZIMUTH_RADIUS * Math.cos(azRad)
                     );
                     elevationHandle.position.set(
-                        -0.8,
-                        ELEVATION_RADIUS * Math.sin(elRad) + CENTER.y,
+                        -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';
+                    let prompt = '💡 Light source from';
                     if (elSnap !== 0) {
                         prompt += ' ' + elevationNames[String(elSnap)];
                     } else {
@@ -619,80 +804,65 @@ class LightingControl3D(gr.HTML):
                     }
                     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);
                 }
-
-                // --------------------------
-                // INTERACTION (MOUSE / TOUCH)
-                // --------------------------
+                
+                // 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
-                    ]);
-
+                    const intersects = raycaster.intersectObjects([azimuthHandle, elevationHandle]);
+                    
                     if (intersects.length > 0) {
                         isDragging = true;
                         dragTarget = intersects[0].object;
-                        dragTarget.material.emissiveIntensity = 1.0;
+                        dragTarget.material.emissiveIntensity = 1.2;
                         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
-                            );
+                            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)
-                                );
+                                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
-                            );
+                            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,
+                                    THREE.MathUtils.radToDeg(Math.atan2(relY, relZ)), 
+                                    -90, 
                                     90
                                 );
                             }
@@ -700,49 +870,46 @@ class LightingControl3D(gr.HTML):
                         updatePositions();
                     } else {
                         raycaster.setFromCamera(mouse, camera);
-                        const intersects = raycaster.intersectObjects([
-                            azimuthHandle,
-                            elevationHandle
-                        ]);
-                        [azimuthHandle, elevationHandle].forEach((h) => {
+                        const intersects = raycaster.intersectObjects([azimuthHandle, elevationHandle]);
+                        [azimuthHandle, elevationHandle].forEach(h => {
                             h.material.emissiveIntensity = 0.6;
                             h.scale.setScalar(1);
                         });
                         if (intersects.length > 0) {
                             intersects[0].object.material.emissiveIntensity = 0.9;
-                            intersects[0].object.scale.setScalar(1.1);
+                            intersects[0].object.scale.setScalar(1.15);
                             canvas.style.cursor = 'grab';
                         } else {
                             canvas.style.cursor = 'default';
                         }
                     }
                 });
-
+                
                 const onMouseUp = () => {
                     if (dragTarget) {
                         dragTarget.material.emissiveIntensity = 0.6;
                         dragTarget.scale.setScalar(1);
-
+                        
+                        // Snap and animate
                         const targetAz = snapToNearest(azimuthAngle, azimuthSteps);
                         const targetEl = snapToNearest(elevationAngle, elevationSteps);
-
-                        const startAz = azimuthAngle;
-                        const startEl = elevationAngle;
+                        
+                        const startAz = azimuthAngle, startEl = elevationAngle;
                         const startTime = Date.now();
-
+                        
                         function 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();
@@ -753,119 +920,90 @@ class LightingControl3D(gr.HTML):
                     dragTarget = null;
                     canvas.style.cursor = 'default';
                 };
-
+                
                 canvas.addEventListener('mouseup', onMouseUp);
                 canvas.addEventListener('mouseleave', onMouseUp);
-
-                // 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;
-
+                
+                // 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.2;
+                        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);
-                        const intersects = raycaster.intersectObjects([
-                            azimuthHandle,
-                            elevationHandle
-                        ]);
-
-                        if (intersects.length > 0) {
-                            isDragging = true;
-                            dragTarget = intersects[0].object;
-                            dragTarget.material.emissiveIntensity = 1.0;
-                            dragTarget.scale.setScalar(1.3);
-                            dragStartMouse.copy(mouse);
-                        }
-                    },
-                    { passive: false }
-                );
-
-                canvas.addEventListener(
-                    'touchmove',
-                    (e) => {
-                        e.preventDefault();
-                        const touch = e.touches[0];
-                        const rect = canvas.getBoundingClientRect();
-                        mouse.x = ((touch.clientX - rect.left) / rect.width) * 2 - 1;
-                        mouse.y = -((touch.clientY - rect.top) / rect.height) * 2 + 1;
-
-                        if (isDragging && dragTarget) {
-                            raycaster.setFromCamera(mouse, camera);
-
-                            if (dragTarget.userData.type === 'azimuth') {
-                                const plane = new THREE.Plane(
-                                    new THREE.Vector3(0, 1, 0),
-                                    -0.05
-                                );
-                                if (raycaster.ray.intersectPlane(plane, intersection)) {
-                                    azimuthAngle = THREE.MathUtils.radToDeg(
-                                        Math.atan2(intersection.x, intersection.z)
-                                    );
-                                    if (azimuthAngle < 0) azimuthAngle += 360;
-                                }
-                            } else if (dragTarget.userData.type === 'elevation') {
-                                const plane = new THREE.Plane(
-                                    new THREE.Vector3(1, 0, 0),
-                                    -0.8
+                        
+                        if (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
                                 );
-                                if (raycaster.ray.intersectPlane(plane, intersection)) {
-                                    const relY = intersection.y - CENTER.y;
-                                    const relZ = intersection.z;
-                                    elevationAngle = THREE.MathUtils.clamp(
-                                        THREE.MathUtils.radToDeg(
-                                            Math.atan2(relY, relZ)
-                                        ),
-                                        -90,
-                                        90
-                                    );
-                                }
                             }
-                            updatePositions();
                         }
-                    },
-                    { passive: false }
-                );
-
-                canvas.addEventListener(
-                    'touchend',
-                    (e) => {
-                        e.preventDefault();
-                        onMouseUp();
-                    },
-                    { passive: false }
-                );
-
-                canvas.addEventListener(
-                    'touchcancel',
-                    (e) => {
-                        e.preventDefault();
-                        onMouseUp();
-                    },
-                    { passive: false }
-                );
-
+                        updatePositions();
+                    }
+                }, { 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();
-
-                // Resize handling
+                
+                // Handle resize
                 new ResizeObserver(() => {
                     camera.aspect = wrapper.clientWidth / wrapper.clientHeight;
                     camera.updateProjectionMatrix();
                     renderer.setSize(wrapper.clientWidth, wrapper.clientHeight);
                 }).observe(wrapper);
-
-                // External updates
+                
+                // Store update functions for external calls
                 wrapper._updateFromProps = (newVal) => {
                     if (newVal && typeof newVal === 'object') {
                         azimuthAngle = newVal.azimuth ?? azimuthAngle;
@@ -873,10 +1011,10 @@ class LightingControl3D(gr.HTML):
                         updatePositions();
                     }
                 };
-
+                
                 wrapper._updateTexture = updateTextureFromUrl;
-
-                // Watch for prop changes (imageUrl and value)
+                
+                // Watch for prop changes
                 let lastImageUrl = props.imageUrl;
                 let lastValue = JSON.stringify(props.value);
                 setInterval(() => {
@@ -895,73 +1033,62 @@ class LightingControl3D(gr.HTML):
                     }
                 }, 100);
             };
-
+            
             initScene();
         })();
         """
-
+        
         super().__init__(
             value=value,
             html_template=html_template,
             js_on_load=js_on_load,
             imageUrl=imageUrl,
-            **kwargs,
+            **kwargs
         )
 
-
-# -----------------------------
-# GRADIO UI
-# -----------------------------
-css = """
+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; }
 #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**. "
-        "Uses the [Multi-Angle-Lighting](https://huggingface.co/dx8152/Qwen-Edit-2509-Multi-Angle-Lighting) LoRA "
-        "with [Rapid-AIO-V19](https://huggingface.co/prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19)."
-    )
-
+    gr.Markdown("Control lighting directions using the **3D viewport** or **sliders**. Use [Multi-Angle-Lighting](https://huggingface.co/dx8152/Qwen-Edit-2509-Multi-Angle-Lighting) LoRA for precise lighting control, paired with [Rapid-AIO-V19](https://huggingface.co/prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19).")
+    
     with gr.Row():
         with gr.Column(scale=1):
             image = gr.Image(label="Input Image", type="pil", height=300)
-
+            
             gr.Markdown("### 3D Lighting Control")
-            gr.Markdown(
-                "*Drag the colored handles: Yellow = Azimuth (direction), Blue = Elevation (height)*"
-            )
-
+            gr.Markdown("*Drag the colored handles: 🟡 Azimuth (Direction), 🔵 Elevation (Height)*")
+            
             lighting_3d = LightingControl3D(
                 value={"azimuth": 0, "elevation": 0},
-                elem_id="lighting-3d-control",
+                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,
                 maximum=315,
                 step=45,
                 value=0,
-                info="0°=front, 90°=right, 180°=rear, 270°=left",
+                info="0°=front, 90°=right, 180°=rear, 270°=left"
             )
-
+            
             elevation_slider = gr.Slider(
                 label="Elevation (Vertical Angle)",
                 minimum=-90,
                 maximum=90,
                 step=90,
                 value=0,
-                info="-90°=from below, 0°=horizontal, 90°=from above",
+                info="-90°=from below, 0°=horizontal, 90°=from above"
             )
 
             with gr.Row():
@@ -971,123 +1098,91 @@ 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
-                )
+                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,
-                )
-                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
-                )
-
-    # Helper callbacks
+                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)
+                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):
-        return build_lighting_prompt(azimuth, elevation)
-
+        """Update prompt preview when sliders change."""
+        prompt = build_lighting_prompt(azimuth, elevation)
+        return prompt
+    
     def sync_3d_to_sliders(lighting_value):
+        """Sync 3D control changes to sliders."""
         if lighting_value and isinstance(lighting_value, dict):
-            az = lighting_value.get("azimuth", 0)
-            el = lighting_value.get("elevation", 0)
+            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
         from io import BytesIO
-
         buffered = BytesIO()
         image.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)
-
-    # Events
+    
     for slider in [azimuth_slider, elevation_slider]:
         slider.change(
             fn=update_prompt_from_sliders,
             inputs=[azimuth_slider, elevation_slider],
-            outputs=[prompt_preview],
+            outputs=[prompt_preview]
         )
-
+    
     lighting_3d.change(
         fn=sync_3d_to_sliders,
         inputs=[lighting_3d],
-        outputs=[azimuth_slider, elevation_slider, prompt_preview],
+        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],
+            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],
+        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],
-        outputs=[width, height],
+        outputs=[width, height]
     ).then(
         fn=update_3d_image,
         inputs=[image],
-        outputs=[lighting_3d],
+        outputs=[lighting_3d]
     )
-
+    
     image.clear(
         fn=lambda: gr.update(imageUrl=None),
-        outputs=[lighting_3d],
+        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,
-        theme=orange_red_theme,
-        mcp_server=True,
-        ssr_mode=False,
-        show_error=True,
-    )
+    css = '.fillable{max-width: 1200px !important}'
+    demo.launch(head=head, css=css, theme=orange_red_theme, mcp_server=True, ssr_mode=False, show_error=True)