Update app.py

app.py

@@ -7,422 +7,363 @@ import json
 import os
 from PIL import Image
 from diffusers import FluxKontextPipeline
-from diffusers.utils import load_image
-from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, list_repo_files
+from diffusers.utils import load_image, peft_utils
+from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard
 from safetensors.torch import load_file
 import requests
 import re
 
-# Load Kontext model
+# Load the base model
 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
-flux_loras_raw = [
-    {
-        "image": "https://huggingface.co/prithivMLmods/FLUX.1-Kontext-Cinematic-Relighting/resolve/main/images/1.png",
-        "title": "Kontext Cinematic Relighting",
-        "repo": "prithivMLmods/FLUX.1-Kontext-Cinematic-Relighting",
-        "trigger_word": "Cinematic Relighting, Relight this portrait with warm, cinematic indoor lighting. Add soft amber highlights and gentle shadows to the face mimicking golden-hour light through a cozy room. Maintain natural skin texture and soft facial shadows, while enhancing eye catchlights for a vivid, lifelike look. Adjust white balance to a warmer tone, and slightly boost exposure to soften the darker midtones. Preserve the subject's pose and expression, and enhance the depth with gentle background bokeh and subtle filmic glow.",
-        "weights": "FLUX.1-Kontext-Cinematic-Relighting.safetensors"
-    },
-]
-print(f"Loaded {len(flux_loras_raw)} LoRAs")
-# Global variables for LoRA management
-current_lora = None
+try:  # Temporary workaround for diffusers LoRA loading issue
+    from diffusers.utils.peft_utils import _derive_exclude_modules
+
+    def new_derive_exclude_modules(*args, **kwargs):
+        exclude_modules = _derive_exclude_modules(*args, **kwargs)
+        if exclude_modules is not None:
+            exclude_modules = [n for n in exclude_modules if "proj_out" not in n]
+        return exclude_modules
+    peft_utils._derive_exclude_modules = new_derive_exclude_modules
+except:
+    pass
+
+# Load LoRA configurations from JSON
+with open("lora_configs.json", "r") as file:
+    data = json.load(file)
+    lora_configs = [
+        {
+            "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"),
+        }
+        for item in data
+    ]
+print(f"Loaded {len(lora_configs)} LoRAs from JSON")
+
+# Global variables for adapter management
+active_lora_adapter = None
 lora_cache = {}
 
 def load_lora_weights(repo_id, weights_filename):
-    """Load LoRA weights from HuggingFace"""
+    """Load adapter weights from HuggingFace"""
     try:
-        # First try with the specified filename
-        try:
+        if repo_id not in lora_cache:
             lora_path = hf_hub_download(repo_id=repo_id, filename=weights_filename)
-            if repo_id not in lora_cache:
-                lora_cache[repo_id] = lora_path
-            return lora_path
-        except Exception as e:
-            print(f"Failed to load {weights_filename}, trying to find alternative LoRA files...")
-            
-            # If the specified file doesn't exist, try to find any .safetensors file
-            from huggingface_hub import list_repo_files
-            try:
-                files = list_repo_files(repo_id)
-                safetensors_files = [f for f in files if f.endswith(('.safetensors', '.bin')) and 'lora' in f.lower()]
-                
-                if not safetensors_files:
-                    # Try without 'lora' in filename
-                    safetensors_files = [f for f in files if f.endswith('.safetensors')]
-                
-                if safetensors_files:
-                    # Try the first available file
-                    for file in safetensors_files:
-                        try:
-                            print(f"Trying alternative file: {file}")
-                            lora_path = hf_hub_download(repo_id=repo_id, filename=file)
-                            if repo_id not in lora_cache:
-                                lora_cache[repo_id] = lora_path
-                            print(f"Successfully loaded alternative LoRA file: {file}")
-                            return lora_path
-                        except:
-                            continue
-                            
-                print(f"No suitable LoRA files found in {repo_id}")
-                return None
-                
-            except Exception as list_error:
-                print(f"Error listing files in repo {repo_id}: {list_error}")
-                return None
-                
+            lora_cache[repo_id] = lora_path
+        return lora_cache[repo_id]
     except Exception as e:
-        print(f"Error loading LoRA from {repo_id}: {e}")
+        print(f"Error loading adapter 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(), None
+def on_lora_select(selected_state: gr.SelectData, lora_configs):
+    """Update UI when an adapter is selected"""
+    if selected_state.index >= len(lora_configs):
+        return "### No adapter selected", gr.update(), None
     
-    lora = flux_loras[selected_state.index]
-    lora_title = lora["title"]
-    lora_repo = lora["repo"]
-    trigger_word = lora["trigger_word"]
+    lora_repo = lora_configs[selected_state.index]["repo"]
+    trigger_word = lora_configs[selected_state.index]["trigger_word"]
     
-    # Create a more informative selected text
-    updated_text = f"### 🎨 Selected Style: {lora_title}"
-    new_placeholder = f"Describe additional details, e.g., 'wearing a red hat' or 'smiling'"
+    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo})"
+    new_placeholder = f"optional description, e.g. 'a man with glasses and a beard'"
     
     return updated_text, gr.update(placeholder=new_placeholder), selected_state.index
 
-def get_huggingface_lora(link):
-    """Download LoRA from HuggingFace link"""
+def fetch_lora_from_hf(link):
+    """Retrieve adapter 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", "")
             
-            # Try to find the correct safetensors file
-            files = list_repo_files(link)
-            safetensors_files = [f for f in files if f.endswith('.safetensors')]
+            fs = HfFileSystem()
+            list_of_files = fs.ls(link, detail=False)
+            safetensors_file = None
             
-            # Prioritize files with 'lora' in the name
-            lora_files = [f for f in safetensors_files if 'lora' in f.lower()]
-            if lora_files:
-                safetensors_file = lora_files[0]
-            elif safetensors_files:
-                safetensors_file = safetensors_files[0]
-            else:
-                # Try .bin files as fallback
-                bin_files = [f for f in files if f.endswith('.bin') and 'lora' in f.lower()]
-                if bin_files:
-                    safetensors_file = bin_files[0]
-                else:
-                    safetensors_file = "pytorch_lora_weights.safetensors"  # Default fallback
+            for file in list_of_files:
+                if file.endswith(".safetensors") and "lora" in file.lower():
+                    safetensors_file = file.split("/")[-1]
+                    break
             
-            print(f"Found LoRA file: {safetensors_file} in {link}")
-            return split_link[1], safetensors_file, trigger_word
+            if not safetensors_file:
+                safetensors_file = "pytorch_lora_weights.safetensors"
             
+            return split_link[1], safetensors_file, trigger_word
         except Exception as e:
-            print(f"Error in get_huggingface_lora: {e}")
-            # Try basic detection
-            try:
-                files = list_repo_files(link)
-                safetensors_file = next((f for f in files if f.endswith('.safetensors')), "pytorch_lora_weights.safetensors")
-                return split_link[1], safetensors_file, ""
-            except:
-                raise Exception(f"Error loading LoRA: {e}")
+            raise Exception(f"Error loading adapter: {e}")
     else:
         raise Exception("Invalid HuggingFace repository format")
 
-def load_custom_lora(link):
-    """Load custom LoRA from user input"""
+def load_user_lora(link):
+    """Load a user-provided adapter"""
     if not link:
-        return gr.update(visible=False), "", gr.update(visible=False), None, gr.Gallery(selected_index=None), "### 🎨 Select an art style from the gallery", None
+        return gr.update(visible=False), "", gr.update(visible=False), None, gr.Gallery(selected_index=None), "### Click on an adapter in the gallery to select it", None
     
     try:
-        repo_name, weights_file, trigger_word = get_huggingface_lora(link)
+        repo_name, weights_file, trigger_word = fetch_lora_from_hf(link)
         
         card = f'''
-        <div class="custom_lora_card">
-            <div style="display: flex; align-items: center; margin-bottom: 12px;">
-                <span style="font-size: 18px; margin-right: 8px;">✅</span>
-                <strong style="font-size: 16px;">Custom LoRA Loaded!</strong>
-            </div>
-            <div style="background: rgba(255, 255, 255, 0.8); padding: 12px; border-radius: 8px;">
-                <h4 style="margin: 0 0 8px 0; color: #333;">{repo_name}</h4>
-                <small style="color: #666;">{"Trigger: <code style='background: #f0f0f0; padding: 2px 6px; border-radius: 4px;'><b>"+trigger_word+"</b></code>" if trigger_word else "No trigger word found"}</small>
+        <div style="border: 1px solid #ddd; padding: 10px; border-radius: 8px; margin: 10px 0;">
+            <span><strong>Loaded custom adapter:</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 = {
+        user_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 Style: {repo_name}", None
+        return gr.update(visible=True), card, gr.update(visible=True), user_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(), "### 🎨 Select an art style from the gallery", None
+        return gr.update(visible=True), f"Error: {str(e)}", gr.update(visible=False), None, gr.update(), "### Click on an adapter in the gallery to select it", None
 
-def remove_custom_lora():
-    """Remove custom LoRA"""
+def unload_user_lora():
+    """Remove the user-provided adapter"""
     return "", gr.update(visible=False), gr.update(visible=False), None, None
 
-def classify_gallery(flux_loras):
-    """Sort gallery by likes"""
-    try:
-        sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True)
-        gallery_items = []
-        
-        for item in sorted_gallery:
-            if "image" in item and "title" in item:
-                image_path = item["image"]
-                title = item["title"]
-                
-                # Simply use the path as-is for Gradio to handle
-                gallery_items.append((image_path, title))
-                print(f"Added to gallery: {image_path} - {title}")
-        
-        print(f"Total gallery items: {len(gallery_items)}")
-        return gallery_items, sorted_gallery
-    except Exception as e:
-        print(f"Error in classify_gallery: {e}")
-        import traceback
-        traceback.print_exc()
-        return [], []
+def sort_lora_gallery(lora_configs):
+    """Sort the adapter gallery by likes"""
+    sorted_gallery = sorted(lora_configs, key=lambda x: x.get("likes", 0), reverse=True)
+    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
 
-def infer_with_lora_wrapper(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
-    """Wrapper function to handle state serialization"""
-    return infer_with_lora(input_image, prompt, selected_index, custom_lora, seed, randomize_seed, guidance_scale, lora_scale, flux_loras, progress)
+def generate_image_wrapper(input_image, prompt, selected_index, user_lora, seed=42, randomize_seed=False, steps=28, guidance_scale=2.5, lora_scale=1.75, width=960, height=1280, lora_configs=None, progress=gr.Progress(track_tqdm=True)):
+    """Wrapper for image generation to handle state"""
+    return generate_image(input_image, prompt, selected_index, user_lora, seed, randomize_seed, steps, guidance_scale, lora_scale, width, height, lora_configs, progress)
 
 @spaces.GPU
-def infer_with_lora(input_image, prompt, selected_index, 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
-    
-    # Check if input image is provided
-    if input_image is None:
-        gr.Warning("Please upload your portrait photo first! 📸")
-        return None, seed, gr.update(visible=False)
+def generate_image(input_image, prompt, selected_index, user_lora, seed=42, randomize_seed=False, steps=28, guidance_scale=2.5, lora_scale=1.0, width=960, height=1280, lora_configs=None, progress=gr.Progress(track_tqdm=True)):
+    """Generate an image using the selected adapter"""
+    global active_lora_adapter, pipe
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-    # Determine which LoRA to use
+    # Select the adapter to use
     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):
-        lora_to_use = flux_loras[selected_index]
-    # Load LoRA if needed
-    if lora_to_use and lora_to_use != current_lora:
+    if user_lora:
+        lora_to_use = user_lora
+    elif selected_index is not None and lora_configs and selected_index < len(lora_configs):
+        lora_to_use = lora_configs[selected_index]
+    print(f"Loaded {len(lora_configs)} adapters from JSON")
+    
+    # Load the adapter if necessary
+    if lora_to_use and lora_to_use != active_lora_adapter:
         try:
-            # Unload current LoRA
-            if current_lora:
-                pipe.unload_lora_weights()
-                print(f"Unloaded previous LoRA")
+            if active_lora_adapter:
+                pipe разгрузить_веса_lora()
             
-            # Load new LoRA
-            repo_id = lora_to_use.get("repo", "unknown")
-            weights_file = lora_to_use.get("weights", "pytorch_lora_weights.safetensors")
-            print(f"Loading LoRA: {repo_id} with weights: {weights_file}")
-            
-            lora_path = load_lora_weights(repo_id, weights_file)
+            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"Successfully loaded: {lora_path} with scale {lora_scale}")
-                current_lora = lora_to_use
-            else:
-                print(f"Failed to load LoRA from {repo_id}")
-                gr.Warning(f"Failed to load {lora_to_use.get('title', 'style')}. Please try a different art style.")
-                return None, seed, gr.update(visible=False)
+                print(f"loaded: {lora_path} with scale {lora_scale}")
+                active_lora_adapter = lora_to_use
             
         except Exception as e:
-            print(f"Error loading LoRA: {e}")
-            # Continue without LoRA
+            print(f"Error loading adapter: {e}")
     else:
-        if lora_to_use:
-            print(f"Using already loaded LoRA: {lora_to_use.get('repo', 'unknown')}")
-    
-    try:
-        # Convert image to RGB
-        input_image = input_image.convert("RGB")
-    except Exception as e:
-        print(f"Error processing image: {e}")
-        gr.Warning("Error processing the uploaded image. Please try a different photo. 📸")
-        return None, seed, gr.update(visible=False)
-    
-    # Check if LoRA is selected
-    if lora_to_use is None:
-        gr.Warning("Please select an art style from the gallery first! 🎨")
-        return None, seed, gr.update(visible=False)
+        print(f"using already loaded adapter: {lora_to_use}")
     
-    # Add trigger word to prompt
-    trigger_word = lora_to_use.get("trigger_word", "")
-    
-    # Special handling for different art styles
-    if trigger_word == "ghibli":
-        prompt = f"Create a Studio Ghibli anime style portrait of the person in the photo, {prompt}. Maintain the facial identity while transforming into whimsical anime art style."
-    elif trigger_word == "homer":
-        prompt = f"Paint the person in Winslow Homer's American realist style, {prompt}. Keep facial features while applying watercolor and marine art techniques."
-    elif trigger_word == "gogh":
-        prompt = f"Transform the portrait into Van Gogh's post-impressionist style with swirling brushstrokes, {prompt}. Maintain facial identity with expressive colors."
-    elif trigger_word == "Cezanne":
-        prompt = f"Render the person in Paul Cézanne's geometric post-impressionist style, {prompt}. Keep facial structure while applying structured brushwork."
-    elif trigger_word == "Renoir":
-        prompt = f"Paint the portrait in Pierre-Auguste Renoir's impressionist style with soft light, {prompt}. Maintain identity with luminous skin tones."
-    elif trigger_word == "claude monet":
-        prompt = f"Create an impressionist portrait in Claude Monet's style with visible brushstrokes, {prompt}. Keep facial features while using light and color."
-    elif trigger_word == "fantasy":
-        prompt = f"Transform into an epic fantasy character portrait, {prompt}. Maintain facial identity while adding magical and fantastical elements."
-    elif trigger_word == ", How2Draw":
+    input_image = input_image.convert("RGB")
+    # Modify prompt based on trigger word
+    trigger_word = lora_to_use["trigger_word"]
+    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"
+    elif trigger_word == "__ ":
+        prompt = f" {prompt}. Accurately render the toolimpact logo and any tool impact iconography. The toolimpact logo begins with a two-line-tall drop-cap capital letter T with a dot in the center of its top bar."
     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}."
+        prompt = f" {prompt}. convert the style of this photo or image to {trigger_word}. Maintain the facial identity of any persons and the general features of the image!"
     
     try:
         image = pipe(
             image=input_image, 
             prompt=prompt,
             guidance_scale=guidance_scale,
+            num_inference_steps=steps,
             generator=torch.Generator().manual_seed(seed),
+            width=width,
+            height=height,
+            max_area=width * height
         ).images[0]
         
         return image, seed, gr.update(visible=True)
     
     except Exception as e:
-        print(f"Error during inference: {e}")
+        print(f"Error during generation: {e}")
         return None, seed, gr.update(visible=False)
 
-# CSS styling with beautiful gradient pastel design
-css = '''
-#gen_btn{height: 100%}
-#gen_column{align-self: stretch}
-#title{text-align: center}
-#title h1{font-size: 3em; display:inline-flex; align-items:center}
-#title img{width: 100px; margin-right: 0.5em}
-#gallery .grid-wrap{height: 10vh}
-#lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
-.card_internal{display: flex;height: 100px;margin-top: .5em}
-.card_internal img{margin-right: 1em}
-.styler{--form-gap-width: 0px !important}
-#progress{height:30px}
-#progress .generating{display:none}
-.progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
-.progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
-'''
+# CSS styling
+css = """
+#app_container {
+    display: flex;
+    gap: 20px;
+}
+#left_panel {
+    min-width: 400px;
+}
+#lora_info {
+    color: #2563eb;
+    font-weight: bold;
+}
+#edit_prompt {
+    flex-grow: 1;
+}
+#generate_button {
+    background: linear-gradient(45deg, #2563eb, #3b82f6);
+    color: white;
+    border: none;
+    padding: 8px 16px;
+    border-radius: 6px;
+    font-weight: bold;
+}
+.user_lora_card {
+    background: #f8fafc;
+    border: 1px solid #e2e8f0;
+    border-radius: 8px;
+    padding: 12px;
+    margin: 8px 0;
+}
+#lora_gallery{
+    overflow: scroll !important
+}
+"""
 
-# Create Gradio interface
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    gr_flux_loras = gr.State(value=flux_loras_raw)
+# Build the Gradio interface
+with gr.Blocks(css=css) as demo:
+    gr_lora_configs = gr.State(value=lora_configs)
     
     title = gr.HTML(
-        """<h1>Flux Kontext DLC 🎈</h1>""",
+        """<h1>Image Style Transfer using FLUX.1 with Adapters</h1>
+        <p>Edit images using custom style adapters. Fast generation with minimal steps.</p>""",
     )
     
     selected_state = gr.State(value=None)
-    custom_loaded_lora = gr.State(value=None)
+    user_lora = gr.State(value=None)
     
-    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 an image for editing", type="pil", height=260)
+    with gr.Row(elem_id="app_container"):
+        with gr.Column(scale=4, elem_id="left_panel"):
+            with gr.Group(elem_id="lora_selection"):
+                input_image = gr.Image(label="Upload a picture", type="pil", height=300)
                 
                 gallery = gr.Gallery(
-                    label="Choose the Flux Kontext LoRA",
+                    label="Pick an Adapter",
                     allow_preview=False,
                     columns=3,
-                    elem_id="gallery",
+                    elem_id="lora_gallery",
                     show_share_button=False,
                     height=400
                 )
                 
-                custom_model = gr.Textbox(
-                    label="🔗 Or use a custom LoRA from HuggingFace", 
-                    placeholder="e.g., username/lora-name",
+                user_lora_input = gr.Textbox(
+                    label="Or enter a custom HuggingFace adapter", 
+                    placeholder="e.g., username/adapter-name",
                     visible=True
                 )
-                custom_model_card = gr.HTML(visible=False)
-                custom_model_button = gr.Button("Remove custom LoRA", visible=False)
+                user_lora_card = gr.HTML(visible=False)
+                unload_user_lora_button = gr.Button("Remove custom adapter", visible=True)
         
         with gr.Column(scale=5):
             with gr.Row():
                 prompt = gr.Textbox(
-                    label="Additional Details (optional)",
+                    label="Editing Prompt",
                     show_label=False,
                     lines=1,
                     max_lines=1,
-                    placeholder="Describe additional details, e.g., 'wearing a red hat' or 'smiling'",
-                    elem_id="prompt"
+                    placeholder="optional description, e.g. 'colorize and stylize, leave all else as is'",
+                    elem_id="edit_prompt"
                 )
-                run_button = gr.Button("Edit Image", elem_id="run_button")
+                run_button = gr.Button("Generate", elem_id="generate_button")
             
-            result = gr.Image(label="Your Kontext Edited Image", interactive=False)
+            result = gr.Image(label="Generated Image", interactive=False)
             reuse_button = gr.Button("Reuse this image", visible=False)
             
-            with gr.Accordion("Advanced Settings", open=False):
+            with gr.Accordion("Advanced Settings", open=True):
                 lora_scale = gr.Slider(
-                    label="Style Strength",
+                    label="Adapter Scale",
                     minimum=0,
                     maximum=2,
                     step=0.1,
-                    value=1.0,
-                    info="How strongly to apply the art style (1.0 = balanced)"
+                    value=1.5,
+                    info="Controls the strength of the adapter effect"
                 )
                 seed = gr.Slider(
-                    label="Random Seed",
+                    label="Seed",
                     minimum=0,
                     maximum=MAX_SEED,
                     step=1,
                     value=0,
-                    info="Set to 0 for random results"
                 )
-                randomize_seed = gr.Checkbox(label="Randomize seed for each generation", value=True)
+                steps = gr.Slider(
+                    label="Steps",
+                    minimum=1,
+                    maximum=40,
+                    value=10,
+                    step=1
+                )
+                width = gr.Slider(
+                    label="Width",
+                    minimum=128,
+                    maximum=2560,
+                    step=1,
+                    value=960,
+                )
+                height = gr.Slider(
+                    label="Height",
+                    minimum=128,
+                    maximum=2560,
+                    step=1,
+                    value=1280,
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                 guidance_scale = gr.Slider(
-                    label="Image Guidance",
+                    label="Guidance Scale",
                     minimum=1,
                     maximum=10,
                     step=0.1,
-                    value=2.5,
-                    info="How closely to follow the input image (lower = more creative)"
+                    value=2.8,
                 )
             
             prompt_title = gr.Markdown(
-                value="### Select an art style from the gallery",
+                value="### Click on an adapter in the gallery to select it",
                 visible=True,
-                elem_id="selected_lora",
+                elem_id="lora_info",
             )
 
     # 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],
+    user_lora_input.input(
+        fn=load_user_lora,
+        inputs=[user_lora_input],
+        outputs=[user_lora_card, user_lora_card, unload_user_lora_button, user_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]
+    unload_user_lora_button.click(
+        fn=unload_user_lora,
+        outputs=[user_lora_input, unload_user_lora_button, user_lora_card, user_lora, selected_state]
     )
     
     gallery.select(
-        fn=update_selection,
-        inputs=[gr_flux_loras],
+        fn=on_lora_select,
+        inputs=[gr_lora_configs],
         outputs=[prompt_title, prompt, selected_state],
         show_progress=False
     )
     
     gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn=infer_with_lora_wrapper,
-        inputs=[input_image, prompt, selected_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, gr_flux_loras],
+        fn=generate_image_wrapper,
+        inputs=[input_image, prompt, selected_state, user_lora, seed, randomize_seed, steps, guidance_scale, lora_scale, width, height, gr_lora_configs],
         outputs=[result, seed, reuse_button]
     )
     
@@ -431,11 +372,12 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
         inputs=[result],
         outputs=[input_image]
     )
-
+    
+    # Initialize the gallery
     demo.load(
-        fn=classify_gallery, 
-        inputs=[gr_flux_loras], 
-        outputs=[gallery, gr_flux_loras]
+        fn=sort_lora_gallery, 
+        inputs=[gr_lora_configs], 
+        outputs=[gallery, gr_lora_configs]
     )
 
 demo.queue(default_concurrency_limit=None)