add lora gallery

app.py

@@ -1,116 +1,379 @@
 import gradio as gr
 import numpy as np
-
 import spaces
 import torch
 import random
+import json
+import os
 from PIL import Image
-
-#from kontext_pipeline import FluxKontextPipeline
-from pipeline_flux_kontext import FluxKontextPipeline
+from kontext_pipeline import FluxKontextPipeline
 from diffusers import FluxTransformer2DModel
 from diffusers.utils import load_image
+from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard
+from safetensors.torch import load_file
+import requests
+import re
 
-from huggingface_hub import hf_hub_download
-
-
+# Load Kontext model
 kontext_path = hf_hub_download(repo_id="diffusers/kontext-v2", filename="dev-opt-2-a-3.safetensors")
-
 MAX_SEED = np.iinfo(np.int32).max
 
 transformer = FluxTransformer2DModel.from_single_file(kontext_path, torch_dtype=torch.bfloat16)
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16).to("cuda")
 
+# Load LoRA data (you'll need to create this JSON file or modify to load your LoRAs)
+try:
+    with open("flux_loras.json", "r") as file:
+        data = json.load(file)
+        flux_loras_raw = [
+            {
+                "image": item["image"],
+                "title": item["title"],
+                "repo": item["repo"],
+                "trigger_word": item.get("trigger_word", ""),
+                "weights": item.get("weights", "pytorch_lora_weights.safetensors"),
+                "likes": item.get("likes", 0),
+                "downloads": item.get("downloads", 0),
+            }
+            for item in data
+        ]
+except FileNotFoundError:
+    # Default LoRAs if JSON file doesn't exist
+    flux_loras_raw = [
+        {
+            "image": "https://via.placeholder.com/300x300?text=LoRA+1",
+            "title": "Example LoRA 1",
+            "repo": "example/lora1",
+            "trigger_word": "style1",
+            "weights": "pytorch_lora_weights.safetensors",
+            "likes": 100,
+            "downloads": 500,
+        },
+        {
+            "image": "https://via.placeholder.com/300x300?text=LoRA+2", 
+            "title": "Example LoRA 2",
+            "repo": "example/lora2",
+            "trigger_word": "style2",
+            "weights": "pytorch_lora_weights.safetensors",
+            "likes": 80,
+            "downloads": 300,
+        }
+    ]
+
+# Global variables for LoRA management
+current_lora = None
+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
+
+def update_selection(selected_state: gr.SelectData, flux_loras):
+    """Update UI when a LoRA is selected"""
+    if selected_state.index >= len(flux_loras):
+        return "### No LoRA selected", gr.update(), selected_state
+    
+    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})"
+    new_placeholder = f"Enter your editing prompt{f' (use {trigger_word} for best results)' if trigger_word else ''}"
+    
+    return updated_text, gr.update(placeholder=new_placeholder), selected_state
+
+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"
+    
+    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}"
+    
+    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"
+
+def remove_custom_lora():
+    """Remove custom LoRA"""
+    return "", gr.update(visible=False), gr.update(visible=False), 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
+
 @spaces.GPU
-def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
+def infer_with_lora(input_image, prompt, selected_state, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
+    """Generate image with selected LoRA"""
+    global current_lora, pipe
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-        
+    
+    # Determine which LoRA to use
+    lora_to_use = None
+    if custom_lora:
+        lora_to_use = custom_lora
+    elif selected_state and flux_loras:
+        selected_index = selected_state.index if hasattr(selected_state, 'index') else None
+        if selected_index is not None and selected_index < len(flux_loras):
+            lora_to_use = flux_loras[selected_index]
+    
+    # Load LoRA if needed
+    if lora_to_use and lora_to_use != current_lora:
+        try:
+            # Unload current LoRA
+            if current_lora:
+                pipe.unload_lora_weights()
+            
+            # Load new LoRA
+            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")
+                current_lora = lora_to_use
+                
+                # Add trigger word to prompt if available
+                trigger_word = lora_to_use.get("trigger_word", "")
+                if trigger_word and trigger_word not in prompt:
+                    prompt = f"{trigger_word} {prompt}"
+            
+        except Exception as e:
+            print(f"Error loading LoRA: {e}")
+            # Continue without LoRA
+    
+    # Set LoRA scale if LoRA is loaded
+    if current_lora and hasattr(pipe, 'set_adapters'):
+        try:
+            pipe.set_adapters("selected_lora", adapter_weights=[lora_scale])
+        except:
+            # Fallback for older diffusers versions
+            pass
+    
     input_image = input_image.convert("RGB")
-    # original_width, original_height = input_image.size
-    
-    # if original_width >= original_height:
-    #     new_width = 1024
-    #     new_height = int(original_height * (new_width / original_width))
-    #     new_height = round(new_height / 64) * 64
-    # else:
-    #     new_height = 1024
-    #     new_width = int(original_width * (new_height / original_height))
-    #     new_width = round(new_width / 64) * 64
-    
-    #input_image_resized = input_image.resize((new_width, new_height), Image.LANCZOS)
-    image = pipe(
-        image=input_image, 
-        prompt=prompt,
-        guidance_scale=guidance_scale,
-        # width=new_width,
-        # height=new_height,
-        generator=torch.Generator().manual_seed(seed),
-    ).images[0]
-    return image, seed, gr.update(visible=True)
+    
+    try:
+        image = pipe(
+            image=input_image, 
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            generator=torch.Generator().manual_seed(seed),
+        ).images[0]
+        
+        return image, seed, gr.update(visible=True)
+    
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        return None, seed, gr.update(visible=False)
 
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 960px;
+# CSS styling
+css = """
+#main_app {
+    display: flex;
+    gap: 20px;
+}
+#box_column {
+    min-width: 400px;
+}
+#gallery_box {
+    border: 1px solid #ddd;
+    border-radius: 8px;
+    padding: 15px;
+}
+#gallery {
+    height: 400px;
+}
+#selected_lora {
+    color: #2563eb;
+    font-weight: bold;
+}
+#prompt {
+    flex-grow: 1;
+}
+#run_button {
+    background: linear-gradient(45deg, #2563eb, #3b82f6);
+    color: white;
+    border: none;
+    padding: 8px 16px;
+    border-radius: 6px;
+    font-weight: bold;
+}
+.custom_lora_card {
+    background: #f8fafc;
+    border: 1px solid #e2e8f0;
+    border-radius: 8px;
+    padding: 12px;
+    margin: 8px 0;
 }
 """
 
+# Create Gradio interface
 with gr.Blocks(css=css) as demo:
+    gr_flux_loras = gr.State(value=flux_loras_raw)
     
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# FLUX.1 Kontext [dev]
-        """)
-
-        with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(label="Upload the image for editing", type="pil")
-                with gr.Row():
-                    prompt = gr.Text(
-                        label="Prompt",
-                        show_label=False,
-                        max_lines=1,
-                        placeholder="Enter your prompt for editing (e.g., 'Remove glasses', 'Add a hat')",
-                        container=False,
-                    )
-                    run_button = gr.Button("Run", scale=0)
-                with gr.Accordion("Advanced Settings", open=False):
-            
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=0,
-                    )
-                    
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                    
-                    guidance_scale = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=1,
-                        maximum=10,
-                        step=0.1,
-                        value=2.5,
-                    )       
-                    
-            with gr.Column():
-                result = gr.Image(label="Result", show_label=False, interactive=False)
-                reuse_button = gr.Button("Reuse this image", visible=False)
-        
+    title = gr.HTML(
+        """<h1> FLUX.1 Kontext Portrait 👩🏻‍🎤
+        <br><small style="font-size: 13px; opacity: 0.75;"></small></h1>""",
+    )
+    
+    selected_state = gr.State()
+    custom_loaded_lora = gr.State()
+    
+    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 image for editing", type="pil", height=250)
+                
+                gallery = gr.Gallery(
+                    label="Pick a LoRA style from the gallery",
+                    allow_preview=False,
+                    columns=3,
+                    elem_id="gallery",
+                    show_share_button=False,
+                    height=400
+                )
+                
+                custom_model = gr.Textbox(
+                    label="Or enter a custom HuggingFace FLUX LoRA", 
+                    placeholder="e.g., username/lora-name",
+                    visible=False
+                )
+                custom_model_card = gr.HTML(visible=False)
+                custom_model_button = gr.Button("Remove custom LoRA", visible=False)
         
+        with gr.Column(scale=5):
+            with gr.Row():
+                prompt = gr.Textbox(
+                    label="Editing Prompt",
+                    show_label=False,
+                    lines=1,
+                    max_lines=1,
+                    placeholder="Enter your editing prompt (e.g., 'Remove glasses', 'Add a hat')",
+                    elem_id="prompt"
+                )
+                run_button = gr.Button("Generate", elem_id="run_button")
+            
+            result = gr.Image(label="Generated Image", interactive=False)
+            reuse_button = gr.Button("Reuse this image", visible=False)
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                lora_scale = gr.Slider(
+                    label="LoRA Scale",
+                    minimum=0,
+                    maximum=2,
+                    step=0.1,
+                    value=1.0,
+                    info="Controls the strength of the LoRA effect"
+                )
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=0,
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                guidance_scale = gr.Slider(
+                    label="Guidance Scale",
+                    minimum=1,
+                    maximum=10,
+                    step=0.1,
+                    value=2.5,
+                )
+            
+            prompt_title = gr.Markdown(
+                value="### Click on a LoRA in the gallery to select it",
+                visible=True,
+                elem_id="selected_lora",
+            )
 
+    # 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],
+    )
+    
+    custom_model_button.click(
+        fn=remove_custom_lora,
+        outputs=[custom_model, custom_model_button, custom_model_card, custom_loaded_lora]
+    )
+    
+    gallery.select(
+        fn=update_selection,
+        inputs=[gr_flux_loras],
+        outputs=[prompt_title, prompt, selected_state],
+        show_progress=False
+    )
+    
     gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn = infer,
-        inputs = [input_image, prompt, seed, randomize_seed, guidance_scale],
-        outputs = [result, seed, reuse_button]
+        fn=infer_with_lora,
+        inputs=[input_image, prompt, selected_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, gr_flux_loras],
+        outputs=[result, seed, reuse_button]
     )
+    
     reuse_button.click(
-        fn = lambda image: image,
-        inputs = [result],
-        outputs = [input_image]
+        fn=lambda image: image,
+        inputs=[result],
+        outputs=[input_image]
+    )
+    
+    # Initialize gallery
+    demo.load(
+        fn=classify_gallery, 
+        inputs=[gr_flux_loras], 
+        outputs=[gallery, gr_flux_loras]
     )
 
+demo.queue(default_concurrency_limit=None)
 demo.launch()