Update app.py

app.py

@@ -125,19 +125,166 @@ DIRECTION_OPTIONS = {
 
 @spaces.GPU
 def infer(input_image, prompt, illumination_dropdown, direction_dropdown, seed=42, randomize_seed=False, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
+    """
+    Performs relighting on an input image using the FLUX.1-Kontext model.
     
+    This function constructs a detailed prompt based on user selections from dropdowns
+    or a custom text prompt. It then uses the diffusers pipeline to generate a new
+    image with the specified lighting, while aiming to preserve the original subject.
+    
+    The logic prioritizes the custom text in the 'prompt' box. If the user types
+    their own prompt, it will be used directly, overriding the 'illumination_dropdown'
+    selection. The dropdowns serve as convenient presets to populate the prompt box.
+    
+    Args:
+        input_image (PIL.Image.Image): The input image to be relighted.
+        prompt (str): The detailed text description of the desired lighting effect.
+                      If this is manually filled, it overrides the illumination dropdown.
+        illumination_dropdown (str): A preset lighting style. See "Dropdown Options" below.
+        direction_dropdown (str): A preset for the light's direction. See "Dropdown Options" below.
+        seed (int): The seed for the random number generator for reproducibility.
+        randomize_seed (bool): If True, a random seed is used, overriding the 'seed' value.
+        guidance_scale (float): Controls how closely the model follows the prompt.
+        progress (gr.Progress): A Gradio progress tracker for the UI.
+    
+    Returns:
+        tuple[PIL.Image.Image, int]: A tuple containing:
+            - The relighted output image.
+            - The seed used for the generation.
+    
+    ------------------------------------------------------------------------------------
+    Dropdown Options
+    ------------------------------------------------------------------------------------
+    
+    **illumination_dropdown Options:**
+    
+      - "natural lighting"
+      - "sunshine from window"
+      - "golden time"
+      - "sunrise in the mountains"
+      - "afternoon light filtering through trees"
+      - "early morning rays, forest clearing"
+      - "golden sunlight streaming through trees"
+      - "sunset over sea"
+      - "golden hour in a meadow"
+      - "golden hour on a city skyline"
+      - "evening glow in the desert"
+      - "dusky evening on a beach"
+      - "mellow evening glow on a lake"
+      - "warm sunset in a rural village"
+      - "moonlight through curtains"
+      - "moonlight in a dark alley"
+      - "midnight in the forest"
+      - "midnight sky with bright starlight"
+      - "fireflies lighting up a summer night"
+      - "warm atmosphere, at home, bedroom"
+      - "home atmosphere, cozy bedroom illumination"
+      - "cozy candlelight"
+      - "candle-lit room, rustic vibe"
+      - "night, cozy warm light from fireplace"
+      - "campfire light"
+      - "neon night, city"
+      - "blue neon light, urban street"
+      - "neon, Wong Kar-wai, warm"
+      - "red and blue police lights in rain"
+      - "red glow, emergency lights"
+      - "sci-fi RGB glowing, cyberpunk"
+      - "rainbow reflections, neon"
+      - "magic lit"
+      - "mystical glow, enchanted forest"
+      - "ethereal glow, magical forest"
+      - "underwater glow, deep sea"
+      - "underwater luminescence"
+      - "aurora borealis glow, arctic landscape"
+      - "crystal reflections in a cave"
+      - "foggy forest at dawn"
+      - "foggy morning, muted light"
+      - "soft, diffused foggy glow"
+      - "stormy sky lighting"
+      - "lightning flash in storm"
+      - "rain-soaked reflections in city lights"
+      - "gentle snowfall at dusk"
+      - "hazy light of a winter morning"
+      - "mysterious twilight, heavy mist"
+      - "vibrant autumn lighting in a forest"
+      - "purple and pink hues at twilight"
+      - "desert sunset with mirage-like glow"
+      - "sunrise through foggy mountains"
+      - "soft studio lighting"
+      - "harsh, industrial lighting"
+      - "fluorescent office lighting"
+      - "harsh spotlight in dark room"
+      - "light and shadow"
+      - "shadow from window"
+      - "apocalyptic, smoky atmosphere"
+      - "evil, gothic, in a cave"
+      - "flickering light in a haunted house"
+      - "golden beams piercing through storm clouds"
+      - "dim candlelight in a gothic castle"
+      - "colorful lantern light at festival"
+      - "golden glow at a fairground"
+      - "soft glow through stained glass"
+      - "glowing embers from a forge"
+      - "custom"
+    
+    **direction_dropdown Options:**
+    - "auto"
+    - "left side"
+    - "right side"
+    - "top"
+    - "top left"
+    - "top right"
+    - "bottom"
+    - "front"
+    - "back"
+    
+    ------------------------------------------------------------------------------------
+    How to Write Custom Prompts
+    ------------------------------------------------------------------------------------
+    
+    For full creative control, write your own prompt in the "Prompt" textbox. This will
+    override any selection from the "Choose Lighting Style" dropdown.
+    
+    **1. Be Descriptive and Specific:**
+    - *Weak Prompt:* `candlelight`
+    - *Strong Prompt:* `Warm orange flickering light from a single candle, casting dramatic,
+      dancing shadows on the wall.`
+    
+    **2. Focus on the Characteristics of Light:**
+    Build your prompt by considering these key elements:
+    - *Color Temperature:* Is the light warm, cool, golden, blue, orange-red, or neutral white?
+    - *Quality & Hardness:* Is it soft and diffused (gentle shadows) or hard and direct
+      (sharp, high-contrast shadows)?
+    - *Source:* Where is the light coming from? A window, a fireplace, a neon sign, the moon?
+    - *Atmospheric Effects:* Do you want `god rays`, `lens flare`, `volumetric fog`, `bokeh`,
+      `dust particles in light beams`, or `smoke`?
+    - *Shadows:* Describe them. Are they `long`, `short`, `deep`, `soft`, or `sharp-edged`?
+    """
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
     input_image = input_image.convert("RGB")
-    
-    # Combine illumination and direction if direction is not "auto"
-    if direction_dropdown != "auto":
-        prompt_prefix = f"with {illumination_dropdown} coming from the {direction_dropdown} of the image"
+
+    #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 == "":
+        prompt = ILLUMINATION_OPTIONS[illumination_dropdown]
+
+    #If the prompt matches the illumination options, prefix that
+    if prompt == ILLUMINATION_OPTIONS[illumination_dropdown]:
+        prompt_prefix = f", with {illumination_dropdown}"
+    #If the prompt was changed, the prefix is empty as the user prompt is predominant
     else:
-        prompt_prefix = f"with {illumination_dropdown} coming from the {direction_dropdown} of the image"
+        prompt_prefix = ""
+
+    # If direction isn't auto, add the direction suffix
+    if direction_dropdown != "auto" and prompt_prefix != "":
+        prompt_prefix = prompt_prefix + f"coming from the {direction_dropdown} of the image"
+    elif direction_dropdown != "auto" and prompt_prefix == "":
+        prompt_prefix = f", light coming from the {direction_dropdown} of the image"
+    
+    prompt_with_template = f"Relight the image{prompt_prefix}. {prompt} Maintain the identity of the foreground subjects."
     
-    prompt_with_template = f"Relight the image, {prompt_prefix}. {prompt} Maintain the identity of the foreground subjects."
+    print(prompt_with_template)
     
     image = pipe(
         image=input_image, 
@@ -191,13 +338,13 @@ with gr.Blocks(css=css) as demo:
                         scale=1
                     )
                 prompt = gr.Textbox(
-                label="Prompt",
-                info="Prompt is generated by the selected illumination style, but you can override it by writing your own",
-                show_label=True,
-                max_lines=3,
-                placeholder="select an illumination style above or type your custom description...",
-                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
+                    label="Prompt",
+                    info="Prompt is generated by the selected illumination style, but you can override it by writing your own",
+                    show_label=True,
+                    max_lines=3,
+                    placeholder="select an illumination style above or type your custom description...",
+                    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
                 )
                 
                 
@@ -256,4 +403,4 @@ with gr.Blocks(css=css) as demo:
         )
     
 
-demo.launch()
+demo.launch(mcp_server=True)