Update app.py

app.py

@@ -132,7 +132,6 @@ Return only the rewritten instruction text directly, without JSON formatting or
     
     return polish_prompt_hf(full_prompt, SYSTEM_PROMPT)
 
-
 # --- Model Loading ---
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -180,11 +179,119 @@ except Exception as e:
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max
 
+# Illumination options mapping
+ILLUMINATION_OPTIONS = {
+# Natural Daylight
+    "natural lighting": "Neutral white color temperature with balanced exposure and soft shadows",
+    "sunshine from window": "Bright directional sunlight with hard shadows and visible light rays",
+    "golden time": "Warm golden hour lighting with enhanced warm colors and soft shadows",
+    "sunrise in the mountains": "Warm backlighting with atmospheric haze and lens flare",
+    "afternoon light filtering through trees": "Dappled sunlight patterns with green color cast from foliage",
+    "early morning rays, forest clearing": "God rays through trees with warm color temperature",
+    "golden sunlight streaming through trees": "Golden god rays with atmospheric particles in light beams",
+    
+    # Sunset & Evening
+    "sunset over sea": "Warm sunset light with soft diffused lighting and gentle gradients",
+    "golden hour in a meadow": "Golden backlighting with lens flare and rim lighting",
+    "golden hour on a city skyline": "Golden lighting on buildings with silhouette effects",
+    "evening glow in the desert": "Warm directional lighting with long shadows",
+    "dusky evening on a beach": "Cool backlighting with horizon silhouettes",
+    "mellow evening glow on a lake": "Warm lighting with water reflections",
+    "warm sunset in a rural village": "Golden hour lighting with peaceful warm tones",
+    
+    # Night & Moonlight
+    "moonlight through curtains": "Cool blue lighting with curtain shadow patterns",
+    "moonlight in a dark alley": "Cool blue lighting with deep urban shadows",
+    "midnight in the forest": "Very low brightness with minimal ambient lighting",
+    "midnight sky with bright starlight": "Cool blue lighting with star point sources",
+    "fireflies lighting up a summer night": "Small glowing points with warm ambient lighting",
+    
+    # Indoor & Cozy
+    "warm atmosphere, at home, bedroom": "Very warm lighting with soft diffused glow",
+    "home atmosphere, cozy bedroom illumination": "Warm table lamp lighting with pools of light",
+    "cozy candlelight": "Warm orange flickering light with dramatic shadows",
+    "candle-lit room, rustic vibe": "Multiple warm candlelight sources with atmospheric shadows",
+    "night, cozy warm light from fireplace": "Warm orange-red firelight with flickering effects",
+    "campfire light": "Warm orange flickering light from below with dancing shadows",
+    
+    # Urban & Neon
+    "neon night, city": "Vibrant blue, magenta, and green neon lights with reflections",
+    "blue neon light, urban street": "Blue neon lighting with urban glow effects",
+    "neon, Wong Kar-wai, warm": "Warm amber and red neon with moody selective lighting",
+    "red and blue police lights in rain": "Alternating red and blue strobing with wet reflections",
+    "red glow, emergency lights": "Red emergency lighting with harsh shadows and high contrast",
+
+     # Sci-Fi & Fantasy
+    "sci-fi RGB glowing, cyberpunk": "Electric blue, pink, and green RGB lighting with glowing effects",
+    "rainbow reflections, neon": "Chromatic rainbow patterns with prismatic reflections",
+    "magic lit": "Colored rim lighting in purple and blue with soft ethereal glow",
+    "mystical glow, enchanted forest": "Supernatural green and blue glowing with floating particles",
+    "ethereal glow, magical forest": "Supernatural lighting with blue-green rim lighting",
+    "underwater glow, deep sea": "Blue-green lighting with caustic patterns and particles",
+    "underwater luminescence": "Blue-green bioluminescent glow with caustic light patterns",
+    "aurora borealis glow, arctic landscape": "Green and purple dancing sky lighting",
+    "crystal reflections in a cave": "Sparkle effects with prismatic light dispersion",
+    
+    # Weather & Atmosphere
+    "foggy forest at dawn": "Volumetric fog with cool god rays through trees",
+    "foggy morning, muted light": "Soft fog effects with reduced contrast throughout",
+    "soft, diffused foggy glow": "Heavy fog with soft lighting and no harsh shadows",
+    "stormy sky lighting": "Dramatic lighting with high contrast and rim lighting",
+    "lightning flash in storm": "Brief intense white light with stark shadows",
+    "rain-soaked reflections in city lights": "Wet surface reflections with streaking light effects",
+    "gentle snowfall at dusk": "Cool blue lighting with snowflake particle effects",
+    "hazy light of a winter morning": "Neutral lighting with atmospheric haze",
+    "mysterious twilight, heavy mist": "Heavy fog with cool lighting and atmospheric depth",
+    
+    # Seasonal & Nature
+    "vibrant autumn lighting in a forest": "Enhanced warm autumn colors with dappled sunlight",
+    "purple and pink hues at twilight": "Warm lighting with soft purple and pink color grading",
+    "desert sunset with mirage-like glow": "Warm orange lighting with heat distortion effects",
+    "sunrise through foggy mountains": "Warm lighting through mist with atmospheric perspective",
+    
+    # Professional & Studio
+    "soft studio lighting": "Multiple diffused sources with even illumination and minimal shadows",
+    "harsh, industrial lighting": "Bright fluorescent lighting with hard shadows",
+    "fluorescent office lighting": "Cool white overhead lighting with slight green tint",
+    "harsh spotlight in dark room": "Single intense directional light with dramatic shadows",
+    
+    # Special Effects & Drama
+    "light and shadow": "Maximum contrast with sharp shadow boundaries",
+    "shadow from window": "Window frame shadow patterns with geometric shapes",
+    "apocalyptic, smoky atmosphere": "Orange-red fire tint with smoke effects",
+    "evil, gothic, in a cave": "Low brightness with cool lighting and deep shadows",
+    "flickering light in a haunted house": "Unstable flickering with cool and warm mixed lighting",
+    "golden beams piercing through storm clouds": "Dramatic god rays with high contrast",
+    "dim candlelight in a gothic castle": "Warm orange candlelight with stone texture enhancement",
+
+    # Festival & Celebration
+    "colorful lantern light at festival": "Multiple colored lantern sources with bokeh effects",
+    "golden glow at a fairground": "Warm carnival lighting with colorful bulb effects",
+    "soft glow through stained glass": "Colored light filtering with rainbow surface patterns",
+    "glowing embers from a forge": "Orange-red glowing particles with intense heat effects"
+
+    }
+
+# Lighting direction options
+DIRECTION_OPTIONS = {
+    "auto": "",  
+    "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."
+}
+
+
 # --- Main Inference Function ---
 @spaces.GPU(duration=60)
 def infer(
     image,
     prompt,
+    illumination_dropdown, direction_dropdown, 
     seed=42,
     randomize_seed=False,
     true_guidance_scale=1.0,
@@ -208,15 +315,30 @@ def infer(
     print(f"Negative Prompt: '{negative_prompt}'")
     print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}")
     
+    #If the dropdown isn't custom, and the user didn't specify a prompt, fill the prompt with the correct one from the illumination options
+    if illumination_dropdown != "custom" and (prompt == "" or prompt == ILLUMINATION_OPTIONS[illumination_dropdown]):
+        prompt = f"change the lighting. add {ILLUMINATION_OPTIONS[illumination_dropdown]}"
+
+
+    # If direction isn't auto, add the direction suffix
+    if direction_dropdown != "auto":
+        prompt_with_template = prompt+ f" coming from the {direction_dropdown}"
+    else:
+        prompt_with_template= prompt
+    
     if rewrite_prompt:
-        prompt = polish_prompt(prompt, image)
-        print(f"Rewritten Prompt: {prompt}")
+        final_prompt = polish_prompt(prompt_with_template, input_image)
+    else:
+        final_prompt = prompt_with_template
+
+    print(f"Calling pipeline with prompt: '{final_prompt}'")
+    
 
     # Generate the edited image - always generate just 1 image
     try:
         images = pipe(
             image,
-            prompt=prompt,
+            prompt=final_prompt,
             negative_prompt=negative_prompt,
             num_inference_steps=num_inference_steps,
             generator=generator,
@@ -231,6 +353,13 @@ def infer(
         print(f"Error during inference: {e}")
         raise e
 
+def update_prompt_from_dropdown(illumination_option):
+    """Update the prompt textbox based on dropdown selection"""
+    if illumination_option == "custom":
+        return ""  # Clear the prompt for custom input
+    else:
+        return ILLUMINATION_OPTIONS[illumination_option]
+
 # --- Examples and UI Layout ---
 examples = [
     # You can add example pairs of [image_path, prompt] here
@@ -273,6 +402,21 @@ with gr.Blocks(css=css) as demo:
                     show_label=True, 
                     type="pil"
                 )
+                with gr.Row():
+                    
+                    illumination_dropdown = gr.Dropdown(
+                        choices=["custom"] + list(ILLUMINATION_OPTIONS.keys()),
+                        value="sunshine from window",
+                        label="Choose Lighting Style",
+                        scale=2
+                    )
+                    
+                    direction_dropdown = gr.Dropdown(
+                        choices=list(DIRECTION_OPTIONS.keys()),
+                        value="auto",
+                        label="Light Direction",
+                        scale=1
+                    )
             # Changed from Gallery to Image
             result = gr.Image(
                 label="Result", 
@@ -324,13 +468,20 @@ with gr.Blocks(css=css) as demo:
             )
 
         # gr.Examples(examples=examples, inputs=[input_image, prompt], outputs=[result, seed], fn=infer, cache_examples=False)
-
+    # update prompt when dropdown changes
+    illumination_dropdown.change(
+        fn=update_prompt_from_dropdown,
+        inputs=[illumination_dropdown],
+        outputs=[prompt]
+    )
+    
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
             input_image,
             prompt,
+            illumination_dropdown, direction_dropdown, 
             seed,
             randomize_seed,
             true_guidance_scale,