Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import numpy as np
-import spaces
+import spaces # This is a special module for Hugging Face Spaces, not needed for local execution
 import torch
 import random
 import json
@@ -13,10 +13,14 @@ from safetensors.torch import load_file
 import requests
 import re
 
-# Load Kontext model
+# Load Kontext model from your local path
 MAX_SEED = np.iinfo(np.int32).max
 
-pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
+# Use the local path for the base model as in your test.py
+pipe = FluxKontextPipeline.from_pretrained(
+    "/hpc2hdd/home/sfei285/Project/Editing/FLUX.1-Kontext-dev", 
+    torch_dtype=torch.bfloat16
+).to("cuda")
 
 # Load LoRA data from our custom JSON file
 with open("kontext_loras.json", "r") as file:
@@ -28,12 +32,12 @@ with open("kontext_loras.json", "r") as file:
             "title": item["title"],
             "repo": item["repo"],
             "weights": item.get("weights", "pytorch_lora_weights.safetensors"),
+            "prompt": item.get("prompt", f"Turn this image into {item['title']} style."),
             # The following keys are kept for compatibility with the original demo structure,
             # but our simplified logic doesn't heavily rely on them.
-            "trigger_word": item.get("trigger_word", ""),
             "lora_type": item.get("lora_type", "flux"),
             "lora_scale_config": item.get("lora_scale", 1.0), # Default scale set to 1.0
-            "prompt_placeholder": item.get("prompt_placeholder", "Describe the subject..."),
+            "prompt_placeholder": item.get("prompt_placeholder", "You can edit the prompt here..."),
         }
         for item in data
     ]
@@ -44,30 +48,30 @@ def update_selection(selected_state: gr.SelectData, flux_loras):
     if selected_state.index >= len(flux_loras):
         return "### No LoRA selected", gr.update(), None, gr.update()
     
-    lora_repo = flux_loras[selected_state.index]["repo"]
+    selected_lora = flux_loras[selected_state.index]
+    lora_repo = selected_lora["repo"]
+    default_prompt = selected_lora.get("prompt")
     
     updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo})"
-    config_placeholder = flux_loras[selected_state.index]["prompt_placeholder"]
 
-    optimal_scale = flux_loras[selected_state.index].get("lora_scale_config", 1.0)
-    print("Selected Style: ", flux_loras[selected_state.index]['title'])
+    optimal_scale = selected_lora.get("lora_scale_config", 1.0)
+    print("Selected Style: ", selected_lora['title'])
     print("Optimal Scale: ", optimal_scale)
-    return updated_text, gr.update(placeholder=config_placeholder), selected_state.index, optimal_scale
+    return updated_text, gr.update(value=default_prompt), selected_state.index, optimal_scale
 
 # This wrapper is kept for compatibility with the Gradio event triggers
-def infer_with_lora_wrapper(input_image, prompt, selected_index, lora_state, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.75,portrait_mode=False, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
+def infer_with_lora_wrapper(input_image, prompt, selected_index, lora_state, custom_lora, seed=0, guidance_scale=2.5, num_inference_steps=28, lora_scale=1.0, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
     """Wrapper function to handle state serialization"""
     # The 'custom_lora' and 'lora_state' arguments are no longer used but kept in the signature
-    return infer_with_lora(input_image, prompt, selected_index, seed, randomize_seed, guidance_scale, lora_scale, portrait_mode, flux_loras, progress)
+    return infer_with_lora(input_image, prompt, selected_index, seed, guidance_scale, num_inference_steps, lora_scale, flux_loras, progress)
 
-@spaces.GPU
-def infer_with_lora(input_image, prompt, selected_index, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, portrait_mode=False, flux_loras=None,  progress=gr.Progress(track_tqdm=True)):
+@spaces.GPU # This decorator is only for Hugging Face Spaces hardware, not needed for local execution
+def infer_with_lora(input_image, prompt, selected_index, seed=0, guidance_scale=2.5, num_inference_steps=28, lora_scale=1.0, flux_loras=None,  progress=gr.Progress(track_tqdm=True)):
     """Generate image with selected LoRA"""
     global pipe
     
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
+    # The seed is now always taken directly from the input. Randomization has been removed.
+    
     # Unload any previous LoRA to ensure a clean state
     if "selected_lora" in pipe.get_active_adapters():
         pipe.unload_lora_weights()
@@ -89,20 +93,12 @@ def infer_with_lora(input_image, prompt, selected_index, seed=42, randomize_seed
             pipe.set_adapters(["selected_lora"], adapter_weights=[lora_scale])
             print(f"Loaded {lora_to_use['repo']} with scale {lora_scale}")
             
-            # Simplified and direct prompt construction
-            style_name = lora_to_use['title']
-            if prompt:
-                final_prompt = f"Turn this image of {prompt} into {style_name} style."
-            else:
-                final_prompt = f"Turn this image into {style_name} style."
-            print(f"Using prompt: {final_prompt}")
-
         except Exception as e:
             print(f"Error loading LoRA: {e}")
-            final_prompt = prompt # Fallback to user prompt if LoRA fails
-    else:
-        # No LoRA selected, just use the original prompt
-        final_prompt = prompt
+
+    # Use the prompt from the textbox directly.
+    final_prompt = prompt
+    print(f"Using prompt: {final_prompt}")
 
     input_image = input_image.convert("RGB")
     
@@ -113,14 +109,17 @@ def infer_with_lora(input_image, prompt, selected_index, seed=42, randomize_seed
             height=input_image.size[1],
             prompt=final_prompt,
             guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
             generator=torch.Generator().manual_seed(seed)
         ).images[0]
         
-        return image, seed, gr.update(visible=True), lora_scale
+        # The seed value is no longer returned, as it's not being changed.
+        return image, gr.update(visible=True), lora_scale
     
     except Exception as e:
         print(f"Error during inference: {e}")
-        return None, seed, gr.update(visible=False), lora_scale
+        # Return an error state for all outputs
+        return None, gr.update(visible=False), lora_scale
 
 # CSS styling
 css = """
@@ -159,6 +158,10 @@ css = """
 #gallery{
     overflow: scroll !important
 }
+/* Custom CSS to ensure the input image is fully visible */
+#input_image_display div[data-testid="image"] img {
+    object-fit: contain !important;
+}
 """
 
 # Create Gradio interface
@@ -169,7 +172,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
         """<h1><img src="https://huggingface.co/spaces/kontext-community/FLUX.1-Kontext-portrait/resolve/main/dora_kontext.png" alt="LoRA"> Kontext-Style LoRA Explorer</h1>""",
         elem_id="title",
     )
-    gr.Markdown("A demo for the style LoRAs from the [Kontext-Style Collection](https://huggingface.co/Kontext-Style) 🤗")
+    gr.Markdown("A demo for the style LoRAs from the [Kontext-Style](https://huggingface.co/Kontext-Style) 🤗")
     
     selected_state = gr.State(value=None)
     # The following states are no longer used by the simplified logic but kept for component structure
@@ -179,15 +182,20 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
     with gr.Row(elem_id="main_app"):
         with gr.Column(scale=4, elem_id="box_column"):
             with gr.Group(elem_id="gallery_box"):
-                input_image = gr.Image(label="Upload a picture of yourself", type="pil", height=300)
-                portrait_mode = gr.Checkbox(label="portrait mode", value=True)
+                input_image = gr.Image(
+                    label="Upload a picture of yourself", 
+                    type="pil", 
+                    height=300,
+                    elem_id="input_image_display"
+                )
                 gallery = gr.Gallery(
                     label="Pick a LoRA",
                     allow_preview=False,
                     columns=3,
                     elem_id="gallery",
                     show_share_button=False,
-                    height=400
+                    height=400,
+                    object_fit="contain"
                 )
                 
                 custom_model = gr.Textbox(
@@ -219,7 +227,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
                     minimum=0,
                     maximum=2,
                     step=0.1,
-                    value=1.5,
+                    value=1.0,
                     info="Controls the strength of the LoRA effect"
                 )
                 seed = gr.Slider(
@@ -229,7 +237,6 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
                     step=1,
                     value=0,
                 )
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                 guidance_scale = gr.Slider(
                     label="Guidance Scale",
                     minimum=1,
@@ -237,6 +244,14 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
                     step=0.1,
                     value=2.5,
                 )
+                num_inference_steps = gr.Slider(
+                    label="Timesteps",
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=28,
+                    info="Number of inference steps"
+                )
             
             prompt_title = gr.Markdown(
                 value="### Click on a LoRA in the gallery to select it",
@@ -257,8 +272,8 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer_with_lora_wrapper,
-        inputs=[input_image, prompt, selected_state, lora_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, portrait_mode, gr_flux_loras],
-        outputs=[result, seed, reuse_button, lora_state]
+        inputs=[input_image, prompt, selected_state, lora_state, custom_loaded_lora, seed, guidance_scale, num_inference_steps, lora_scale, gr_flux_loras],
+        outputs=[result, reuse_button, lora_state]
     )
     
     reuse_button.click(
@@ -269,7 +284,8 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
     
     # Initialize gallery
     demo.load(
-        fn=lambda: (flux_loras_raw, flux_loras_raw), 
+        fn=lambda loras: ([(item["image"], item["title"]) for item in loras], loras), 
+        inputs=[gr_flux_loras], 
         outputs=[gallery, gr_flux_loras]
     )