Update app.py

app.py

@@ -21,7 +21,6 @@ MAX_SEED = np.iinfo(np.int32).max
 ILLUMINATION_OPTIONS = {
     "sunset over sea": "add warm sunset light reflecting off water, creating a peaceful atmosphere with soft, diffused lighting and gentle color gradients",
     "sunshine from window": "Add directional sunlight from window source. Increase brightness on lit areas. Create hard shadows with sharp edges. Set warm white color temperature. Add visible light rays and dust particles in beams.",
-    "sunset over sea": "Set warm orange-red color temperature. Increase warm tones substantially. Add horizontal light source from background. Create soft gradient lighting. Add water reflections if visible.",
     "golden time": "Set golden color temperature. Increase contrast moderately. Add rim lighting on subjects. Enhance warm colors. Create soft directional shadows. Reduce blue tones.",
     "sci-fi RGB glowing, cyberpunk": "add futuristic RGB lighting with electric blues, hot pinks, and neon greens creating a high-tech atmosphere with dramatic color separation and glowing effects",
     "neon night, city":"add vibrant neon lights in electric blues, magentas, and greens casting colorful reflections on surfaces, creating a cyberpunk urban atmosphere with dramatic color contrasts",
@@ -78,7 +77,6 @@ ILLUMINATION_OPTIONS = {
     "soft, diffused foggy glow": "Add volumetric fog throughout scene. Reduce contrast substantially. Create soft lighting with no harsh shadows. Set neutral color temperature.",
     "underwater luminescence": "Set blue-green color temperature. Add bioluminescent glowing effects. Create caustic light patterns. Add particle effects in water.",
     "rain-soaked reflections in city lights": "Add reflections on wet surfaces. Create light streaking effects. Set urban lighting with multiple colored sources. Add rain particle effects.",
-    "golden time": "add warm golden hour lighting with soft, diffused sunlight creating a magical glow, enhanced contrast, and rich warm tones throughout the scene",
     "golden sunlight streaming through trees": "Add god rays effect. Set golden color temperature. Create dappled light patterns. Add atmospheric particles in light beams.",
     "fireflies lighting up a summer night": "Add small glowing points of light. Set warm color temperature. Create subtle animated glow effects. Maintain dark ambient lighting.",
     "glowing embers from a forge": "Add orange-red glowing particle effects. Set very warm color temperature. Create intense heat glow. Add sparking effects.",
@@ -90,15 +88,34 @@ ILLUMINATION_OPTIONS = {
     "golden beams piercing through storm clouds": "Add dramatic god rays through clouds. Set golden color temperature. Create high contrast between light and shadow areas."
 }
 
+# Lighting direction options
+DIRECTION_OPTIONS = {
+    "none": "",  
+    "left side": "Position the light source from the left side of the frame, creating shadows falling to the right.",
+    "right side": "Position the light source from the right side of the frame, creating shadows falling to the left.",
+    "top": "Position the light source from directly above, creating downward shadows.",
+    "top left": "Position the light source from the top left corner, creating diagonal shadows falling down and to the right.",
+    "top right": "Position the light source from the top right corner, creating diagonal shadows falling down and to the left.",
+    "bottom": "Position the light source from below, creating upward shadows and dramatic under-lighting.",
+    "front": "Position the light source from the front, minimizing shadows and creating even illumination.",
+    "back": "Position the light source from behind the subject, creating silhouette effects and rim lighting."
+}
+
 @spaces.GPU
-def infer(input_image, prompt, illumination_dropdown, seed=42, randomize_seed=False, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
+def infer(input_image, prompt, illumination_dropdown, direction_dropdown, seed=42, randomize_seed=False, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
     input_image = input_image.convert("RGB")
     
-    prompt_with_template = f"Relight the image, {prompt}. Focus on how the new lighting affects shadows, highlights, color temperature, and overall mood of the scene.Maintain the identiity of the forground subjects"
+    # Combine illumination and direction if direction is not "Auto"
+    if direction_dropdown != "none":
+        prompt_prefix = f"with {illumination_dropdown} coming from the {direction_dropdown}"
+    else:
+        prompt_prefix = f"with {illumination_dropdown} coming from the {direction_dropdown}"
+    
+    prompt_with_template = f"Relight the image, {prompt_prefix}. {prompt} Maintain the identity of the foreground subjects."
     
     image = pipe(
         image=input_image, 
@@ -112,7 +129,7 @@ def infer(input_image, prompt, illumination_dropdown, seed=42, randomize_seed=Fa
 
 def update_prompt_from_dropdown(illumination_option):
     """Update the prompt textbox based on dropdown selection"""
-    if illumination_option == "Custom":
+    if illumination_option == "custom":
         return ""  # Clear the prompt for custom input
     else:
         return ILLUMINATION_OPTIONS[illumination_option]
@@ -135,23 +152,32 @@ with gr.Blocks(css=css) as demo:
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(label="Upload the image for relighting", type="pil")
+                
+                with gr.Row():
+                    illumination_dropdown = gr.Dropdown(
+                        choices=list(ILLUMINATION_OPTIONS.keys()) + ["custom"],
+                        value="sunshine from window",
+                        label="Choose Illumination Style",
+                        info="Select a predefined lighting style or choose 'Custom' to use your own prompt",
+                        scale=2
+                    )
+                    
+                    direction_dropdown = gr.Dropdown(
+                        choices=list(DIRECTION_OPTIONS.keys()),
+                        value="auto",
+                        label="Light Direction",
+                        info="Choose the direction of the light source",
+                        scale=1
+                    )
+                
                 prompt = gr.Textbox(
-                    label="prompt",
+                    label="Lighting Description",
                     show_label=True,
                     max_lines=3,
                     placeholder="Select an illumination style above or enter your custom description...",
-                    value="bright natural sunlight streaming through a window, creating warm golden beams and soft shadows across the scene, with high contrast between lit and shadowed areas",
+                    value="Add directional sunlight from window source. Increase brightness on lit areas. Create hard shadows with sharp edges. Set warm white color temperature. Add visible light rays and dust particles in beams.",
                     container=True
                 )
-            
-                illumination_dropdown = gr.Dropdown(
-                    choices=list(ILLUMINATION_OPTIONS.keys()) + ["Custom"],
-                    value="sunshine from window",
-                    label="Choose Illumination Style",
-                    info="Select a predefined lighting style or choose 'Custom' to use your own prompt"
-                )
-                
-                
                 
                 run_button = gr.Button("Run", scale=0, variant="primary")
                 
@@ -187,14 +213,14 @@ with gr.Blocks(css=css) as demo:
         
         gr.Examples(
             examples=[
-                ["./assets/pexels-creationhill-1681010.jpg", "Add multiple colored light sources from lanterns. Create warm festival lighting. Set varied color temperatures. Add bokeh effects.", "colorful lantern light at festival", 0, True, 2.5],
-                ["./assets/pexels-creationhill-1681010.jpg",  "add futuristic RGB lighting with electric blues, hot pinks, and neon greens creating a high-tech atmosphere with dramatic color separation and glowing effects", "sci-fi RGB glowing, cyberpunk",0, True, 2.5],
-                ["./assets/pexels-moose-photos-170195-1587009.jpg",  "Set blue-green color temperature. Add volumetric lighting effects. Reduce red channel significantly. Create particle effects in light beams. Add caustic light patterns.", "underwater glow, deep sea",0, True, 2.5],
-                ["./assets/pexels-moose-photos-170195-1587009.jpg", "Replace lighting with red sources. Add flashing strobing effects. Increase contrast. Create harsh shadows. Set monochromatic red color scheme.", "red glow, emergency lights", 0, True, 2.5],
-                ["./assets/pexels-simon-robben-55958-614810.jpg",  "Add directional sunlight from window source. Increase brightness on lit areas. Create hard shadows with sharp edges. Set warm white color temperature. Add visible light rays and dust particles in beams.", "sunshine from window",0, True, 2.5],
-                ["./assets/pexels-simon-robben-55958-614810.jpg", "add vibrant neon lights in electric blues, magentas, and greens casting colorful reflections on surfaces, creating a cyberpunk urban atmosphere with dramatic color contrasts", "neon night, city", 0, True, 2.5]
+                ["./assets/pexels-creationhill-1681010.jpg", "Add multiple colored light sources from lanterns. Create warm festival lighting. Set varied color temperatures. Add bokeh effects.", "colorful lantern light at festival", "auto", 0, True, 2.5],
+                ["./assets/pexels-creationhill-1681010.jpg",  "add futuristic RGB lighting with electric blues, hot pinks, and neon greens creating a high-tech atmosphere with dramatic color separation and glowing effects", "sci-fi RGB glowing, cyberpunk", "left side", 0, True, 2.5],
+                ["./assets/pexels-moose-photos-170195-1587009.jpg",  "Set blue-green color temperature. Add volumetric lighting effects. Reduce red channel significantly. Create particle effects in light beams. Add caustic light patterns.", "underwater glow, deep sea", "top", 0, True, 2.5],
+                ["./assets/pexels-moose-photos-170195-1587009.jpg", "Replace lighting with red sources. Add flashing strobing effects. Increase contrast. Create harsh shadows. Set monochromatic red color scheme.", "red glow, emergency lights", "right side", 0, True, 2.5],
+                ["./assets/pexels-simon-robben-55958-614810.jpg",  "Add directional sunlight from window source. Increase brightness on lit areas. Create hard shadows with sharp edges. Set warm white color temperature. Add visible light rays and dust particles in beams.", "sunshine from window", "top right", 0, True, 2.5],
+                ["./assets/pexels-simon-robben-55958-614810.jpg", "add vibrant neon lights in electric blues, magentas, and greens casting colorful reflections on surfaces, creating a cyberpunk urban atmosphere with dramatic color contrasts", "neon night, city", "top left", 0, True, 2.5]
             ],
-            inputs=[input_image, prompt, illumination_dropdown, seed, randomize_seed, guidance_scale],
+            inputs=[input_image, prompt, illumination_dropdown, direction_dropdown, seed, randomize_seed, guidance_scale],
             outputs=[result, seed],
             fn=infer,
             cache_examples="lazy"
@@ -203,7 +229,7 @@ with gr.Blocks(css=css) as demo:
         gr.on(
             triggers=[run_button.click, prompt.submit],
             fn = infer,
-            inputs = [input_image, prompt, illumination_dropdown, seed, randomize_seed, guidance_scale],
+            inputs = [input_image, prompt, illumination_dropdown, direction_dropdown, seed, randomize_seed, guidance_scale],
             outputs = [result, seed]
         )