Update app.py

app.py

@@ -18,38 +18,26 @@ MAX_SEED = np.iinfo(np.int32).max
 
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
 
-# Load LoRA data (you'll need to create this JSON file or modify to load your LoRAs)
-
-with open("flux_loras.json", "r") as file:
+# Load LoRA data from our custom JSON file
+with open("kontext_loras.json", "r") as file:
     data = json.load(file)
+    # Add default values for keys that might be missing, to prevent errors
     flux_loras_raw = [
         {
             "image": item["image"],
             "title": item["title"],
             "repo": item["repo"],
-            "trigger_word": item.get("trigger_word", ""),
-            "trigger_position": item.get("trigger_position", "prepend"),
             "weights": item.get("weights", "pytorch_lora_weights.safetensors"),
+            # 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.5),
-            "prompt_placeholder": item.get("prompt_placeholder", ""),
+            "lora_scale_config": item.get("lora_scale", 1.0), # Default scale set to 1.0
+            "prompt_placeholder": item.get("prompt_placeholder", "Describe the subject..."),
         }
         for item in data
     ]
-print(f"Loaded {len(flux_loras_raw)} LoRAs from JSON")
-# Global variables for LoRA management
-lora_cache = {}
-
-def load_lora_weights(repo_id, weights_filename):
-    """Load LoRA weights from HuggingFace"""
-    try:
-        if repo_id not in lora_cache:
-            lora_path = hf_hub_download(repo_id=repo_id, filename=weights_filename)
-            lora_cache[repo_id] = lora_path
-        return lora_cache[repo_id]
-    except Exception as e:
-        print(f"Error loading LoRA from {repo_id}: {e}")
-        return None
+print(f"Loaded {len(flux_loras_raw)} LoRAs from kontext_loras.json")
 
 def update_selection(selected_state: gr.SelectData, flux_loras):
     """Update UI when a LoRA is selected"""
@@ -57,148 +45,73 @@ def update_selection(selected_state: gr.SelectData, flux_loras):
         return "### No LoRA selected", gr.update(), None, gr.update()
     
     lora_repo = flux_loras[selected_state.index]["repo"]
-    trigger_word = flux_loras[selected_state.index]["trigger_word"]
     
     updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo})"
     config_placeholder = flux_loras[selected_state.index]["prompt_placeholder"]
-    if config_placeholder:
-        new_placeholder = config_placeholder
-    else:
-        new_placeholder = f"opt - describe the person/subject, e.g. 'a man with glasses and a beard'"
 
-    print("Selected Index: ", flux_loras[selected_state.index])
-    optimal_scale = flux_loras[selected_state.index].get("lora_scale_config", 1.5)
+    optimal_scale = flux_loras[selected_state.index].get("lora_scale_config", 1.0)
+    print("Selected Style: ", flux_loras[selected_state.index]['title'])
     print("Optimal Scale: ", optimal_scale)
-    return updated_text, gr.update(placeholder=new_placeholder), selected_state.index, optimal_scale
-
-
-def get_huggingface_lora(link):
-    """Download LoRA from HuggingFace link"""
-    split_link = link.split("/")
-    if len(split_link) == 2:
-        try:
-            model_card = ModelCard.load(link)
-            trigger_word = model_card.data.get("instance_prompt", "")
-            
-            fs = HfFileSystem()
-            list_of_files = fs.ls(link, detail=False)
-            safetensors_file = None
-            
-            for file in list_of_files:
-                if file.endswith(".safetensors") and "lora" in file.lower():
-                    safetensors_file = file.split("/")[-1]
-                    break
-            
-            if not safetensors_file:
-                safetensors_file = "pytorch_lora_weights.safetensors"
-            
-            return split_link[1], safetensors_file, trigger_word
-        except Exception as e:
-            raise Exception(f"Error loading LoRA: {e}")
-    else:
-        raise Exception("Invalid HuggingFace repository format")
-
-def load_custom_lora(link):
-    """Load custom LoRA from user input"""
-    if not link:
-        return gr.update(visible=False), "", gr.update(visible=False), None, gr.Gallery(selected_index=None), "### Click on a LoRA in the gallery to select it", None
-    
-    try:
-        repo_name, weights_file, trigger_word = get_huggingface_lora(link)
-        
-        card = f'''
-        <div style="border: 1px solid #ddd; padding: 10px; border-radius: 8px; margin: 10px 0;">
-            <span><strong>Loaded custom LoRA:</strong></span>
-            <div style="margin-top: 8px;">
-                <h4>{repo_name}</h4>
-                <small>{"Using: <code><b>"+trigger_word+"</b></code> as trigger word" if trigger_word else "No trigger word found"}</small>
-            </div>
-        </div>
-        '''
-        
-        custom_lora_data = {
-            "repo": link,
-            "weights": weights_file,
-            "trigger_word": trigger_word
-        }
-        
-        return gr.update(visible=True), card, gr.update(visible=True), custom_lora_data, gr.Gallery(selected_index=None), f"Custom: {repo_name}", None
-    
-    except Exception as e:
-        return gr.update(visible=True), f"Error: {str(e)}", gr.update(visible=False), None, gr.update(), "### Click on a LoRA in the gallery to select it", None
-
-def remove_custom_lora():
-    """Remove custom LoRA"""
-    return "", gr.update(visible=False), gr.update(visible=False), None, None
-
-def classify_gallery(flux_loras):
-    """Sort gallery by likes"""
-    sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True)
-    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
+    return updated_text, gr.update(placeholder=config_placeholder), 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)):
     """Wrapper function to handle state serialization"""
-    return infer_with_lora(input_image, prompt, selected_index, lora_state, custom_lora, seed, randomize_seed, guidance_scale, lora_scale, portrait_mode, flux_loras, progress)
+    # 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)
 
 @spaces.GPU
-def infer_with_lora(input_image, prompt, selected_index, lora_state, custom_lora, 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)):
+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)):
     """Generate image with selected LoRA"""
     global pipe
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+
+    # Unload any previous LoRA to ensure a clean state
+    if "selected_lora" in pipe.get_active_adapters():
+        pipe.unload_lora_weights()
     
-    # Determine which LoRA to use
+    # Determine which LoRA to use from our gallery
     lora_to_use = None
-    if custom_lora:
-        lora_to_use = custom_lora
-    elif selected_index is not None and flux_loras and selected_index < len(flux_loras):
+    if selected_index is not None and flux_loras and selected_index < len(flux_loras):
         lora_to_use = flux_loras[selected_index]
-    print(f"Loaded {len(flux_loras)} LoRAs from JSON")
-    # Load LoRA if needed
-    print(f"LoRA to use: {lora_to_use}")
+    
     if lora_to_use:
+        print(f"Applying LoRA: {lora_to_use['title']}")
         try:
-            if "selected_lora" in pipe.get_active_adapters():
-                pipe.unload_lora_weights()
-            
-            lora_path = load_lora_weights(lora_to_use["repo"], lora_to_use["weights"])
-            if lora_path:
-                pipe.load_lora_weights(lora_path, adapter_name="selected_lora")
-                pipe.set_adapters(["selected_lora"], adapter_weights=[lora_scale])
-                print(f"loaded: {lora_path} with scale {lora_scale}")
+            # Load LoRA directly from the Hugging Face Hub
+            pipe.load_lora_weights(
+                lora_to_use["repo"], 
+                weight_name=lora_to_use["weights"], 
+                adapter_name="selected_lora"
+            )
+            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
+
     input_image = input_image.convert("RGB")
-    # Add trigger word to prompt
-    trigger_word = lora_to_use["trigger_word"]
-    is_kontext_lora = lora_to_use["lora_type"] == "kontext"
-    if not is_kontext_lora:
-        if portrait_mode:
-            if trigger_word == ", How2Draw":
-                prompt = f"create a How2Draw sketch of the person of the photo {prompt}, maintain the facial identity of the person and general features"
-            elif trigger_word == ", video game screenshot in the style of THSMS":
-                prompt = f"create a video game screenshot in the style of THSMS with the person from the photo, {prompt}. maintain the facial identity of the person and general features"
-            else:
-                prompt = f"convert the style of this portrait photo to {trigger_word} while maintaining the identity of the person. {prompt}. Make sure to maintain the person's facial identity and features, while still changing the overall style to {trigger_word}."
-        else:
-           if trigger_word == ", How2Draw":
-                prompt = f"create a How2Draw sketch of the photo {prompt}"
-           elif trigger_word == ", video game screenshot in the style of THSMS":
-                prompt = f"create a video game screenshot in the style of THSMS of the photo, {prompt}."
-           else:
-                prompt = f"convert the style of this photo {prompt} to {trigger_word}." 
-    else: 
-        prompt = f"{trigger_word}. {prompt}."
+    
     try:
         image = pipe(
             image=input_image,
             width=input_image.size[0],
             height=input_image.size[1],
-            prompt=prompt,
+            prompt=final_prompt,
             guidance_scale=guidance_scale,
             generator=torch.Generator().manual_seed(seed)
         ).images[0]
@@ -253,12 +166,13 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
     gr_flux_loras = gr.State(value=flux_loras_raw)
     
     title = gr.HTML(
-        """<h1><img src="https://huggingface.co/spaces/kontext-community/FLUX.1-Kontext-portrait/resolve/main/dora_kontext.png" alt="LoRA"> FLUX.1 Kontext LoRA the Explorer</h1>""",
+        """<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("Flux.1 Kontext [dev] + community Kontext & Flux LoRAs 🤗")
+    gr.Markdown("A demo for the style LoRAs from the [Kontext-Style Collection](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
     custom_loaded_lora = gr.State(value=None)
     lora_state = gr.State(value=1.0)
     
@@ -331,16 +245,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
             )
 
     # Event handlers
-    custom_model.input(
-        fn=load_custom_lora,
-        inputs=[custom_model],
-        outputs=[custom_model_card, custom_model_card, custom_model_button, custom_loaded_lora, gallery, prompt_title, selected_state],
-    )
-    
-    custom_model_button.click(
-        fn=remove_custom_lora,
-        outputs=[custom_model, custom_model_button, custom_model_card, custom_loaded_lora, selected_state]
-    )
+    # The custom model inputs are no longer needed as we've hidden them.
     
     gallery.select(
         fn=update_selection,
@@ -364,8 +269,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
     
     # Initialize gallery
     demo.load(
-        fn=classify_gallery, 
-        inputs=[gr_flux_loras], 
+        fn=lambda: (flux_loras_raw, flux_loras_raw), 
         outputs=[gallery, gr_flux_loras]
     )