update app

app.py

@@ -1,4 +1,5 @@
 import os
+import gc
 import gradio as gr
 import numpy as np
 import spaces
@@ -7,12 +8,10 @@ import random
 from PIL import Image
 from typing import Iterable
 
-# Pipeline for FLUX.2 Klein
 from diffusers import Flux2KleinPipeline
 from diffusers.utils import load_image
 from huggingface_hub import hf_hub_download
 
-# --- Hardware and Theme Setup ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 from gradio.themes import Soft
@@ -62,71 +61,128 @@ class OrangeRedTheme(Soft):
 orange_red_theme = OrangeRedTheme()
 MAX_SEED = np.iinfo(np.int32).max
 
-# --- Model Loading ---
-print("Loading FLUX.2 Klein 9B model...")
+ADAPTER_SPECS = {
+    "Simple-Tuner": {
+        "repo": "markury/flux2k9b-simpletuner-lora-loona",
+        "weights": "pytorch_lora_weights.safetensors",
+        "adapter_name": "simple-tuner"
+    },
+    "Klein-Delight-Style": {
+        "repo": "linoyts/Flux2-Klein-Delight-LoRA",
+        "weights": "pytorch_lora_weights.safetensors",
+        "adapter_name": "klein-delight"
+    }
+}
+
+LOADED_ADAPTERS = set()
+
+print("Loading FLUX.2 Klein 9B model base...")
 pipe = Flux2KleinPipeline.from_pretrained(
     "black-forest-labs/FLUX.2-klein-9B",
     torch_dtype=torch.bfloat16
 ).to(device)
-print("Model loaded successfully.")
-
-# --- LoRA Loading (Updated) ---
-print("Loading new LoRA adapters...")
-pipe.load_lora_weights(
-    "markury/flux2k9b-simpletuner-lora-loona",
-    weight_name="pytorch_lora_weights.safetensors",
-    adapter_name="simple-tuner"
-)
-pipe.load_lora_weights(
-    "linoyts/Flux2-Klein-Delight-LoRA",
-    weight_name="pytorch_lora_weights.safetensors",
-    adapter_name="klein-delight"
-)
-print("All LoRA adapters loaded.")
+print("Base Model loaded successfully.")
 
-# Updated map for the new adapters
-ADAPTER_MAP = {
-    "Simple-Tuner": "simple-tuner",
-    "Klein-Delight-Style": "klein-delight",
-}
+def update_dimensions_on_upload(image):
+    """Resizes image to be divisible by 16 to avoid tensor mismatch errors in FLUX."""
+    if image is None:
+        return 1024, 1024
+    
+    original_width, original_height = image.size
+    
+    # Simple logic to keep aspect ratio but snap to multiples of 16
+    # You can adjust the max size (e.g., 1024) as needed
+    scale = min(1024 / original_width, 1024 / original_height)
+    new_width = int(original_width * scale)
+    new_height = int(original_height * scale)
+    
+    new_width = (new_width // 16) * 16
+    new_height = (new_height // 16) * 16
+    
+    return new_width, new_height
 
 @spaces.GPU
-def infer(input_image, prompt, lora_adapter, seed=42, randomize_seed=True, guidance_scale=4.0, steps=4, progress=gr.Progress(track_tqdm=True)):
-    # Input image is required for image-to-image tasks
+def infer(
+    input_image, 
+    prompt, 
+    lora_adapter, 
+    seed=42, 
+    randomize_seed=True, 
+    guidance_scale=1.0, 
+    steps=4, 
+    progress=gr.Progress(track_tqdm=True)
+):
+    gc.collect()
+    torch.cuda.empty_cache()
+
     if not input_image:
         raise gr.Error("Please upload an image to apply a style to.")
 
-    # Dynamically set the adapter based on the dropdown choice
-    adapter_name = ADAPTER_MAP.get(lora_adapter)
-    if adapter_name:
-        print(f"Activating LoRA: {lora_adapter} ({adapter_name})")
+    spec = ADAPTER_SPECS.get(lora_adapter)
+    if spec:
+        adapter_name = spec["adapter_name"]
+        
+        if adapter_name not in LOADED_ADAPTERS:
+            print(f"--- Downloading and Loading Adapter: {lora_adapter} ---")
+            try:
+                pipe.load_lora_weights(
+                    spec["repo"], 
+                    weight_name=spec["weights"], 
+                    adapter_name=adapter_name
+                )
+                LOADED_ADAPTERS.add(adapter_name)
+            except Exception as e:
+                raise gr.Error(f"Failed to load adapter {lora_adapter}: {e}")
+        else:
+            print(f"--- Adapter {lora_adapter} is already loaded. ---")
+            
+        print(f"Activating LoRA: {adapter_name}")
         pipe.set_adapters([adapter_name], adapter_weights=[1.0])
     else:
-        # If "None" is selected (or an invalid choice), disable LoRAs
-        print("No LoRA selected. Disabling adapters.")
+        print("No valid LoRA selected or found. Disabling adapters.")
         pipe.disable_lora()
-        
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-    original_image = input_image.copy().convert("RGB")
+    width, height = update_dimensions_on_upload(input_image)
+    processed_input = input_image.resize((width, height), Image.LANCZOS).convert("RGB")
     
-    image = pipe(
-        image=original_image, 
-        prompt=prompt,
-        guidance_scale=guidance_scale,
-        width=original_image.size[0],
-        height=original_image.size[1],
-        num_inference_steps=steps,
-        generator=torch.Generator(device=device).manual_seed(seed),
-    ).images[0]
+    # 4. Inference
+    try:
+        image = pipe(
+            image=processed_input, 
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            width=width,
+            height=height,
+            num_inference_steps=steps,
+            generator=torch.Generator(device=device).manual_seed(seed),
+        ).images[0]
+        
+        return image, seed
 
-    return image, seed
+    except Exception as e:
+        raise gr.Error(f"Inference failed: {e}")
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache()
 
 @spaces.GPU
 def infer_example(input_image, prompt, lora_adapter):
     # Use a fixed seed for reproducible examples
-    image, seed = infer(input_image, prompt, lora_adapter, seed=12345, randomize_seed=False)
+    if input_image is None: 
+        return None, 0
+        
+    image, seed = infer(
+        input_image=input_image, 
+        prompt=prompt, 
+        lora_adapter=lora_adapter, 
+        seed=12345, 
+        randomize_seed=False,
+        guidance_scale=4.0,
+        steps=4
+    )
     return image, seed
 
 # --- UI Layout ---
@@ -140,18 +196,26 @@ with gr.Blocks() as demo:
         gr.Markdown("# **FLUX.2 Klein LoRA Stylizer**", elem_id="main-title")
         gr.Markdown(
             "Apply creative styles to your images using **FLUX.2-klein-9B** and specialized LoRA adapters. "
-            "Upload an image, select a style, and write a prompt to guide the transformation."
+            "Adapters are loaded lazily (on-demand) to save resources."
         )
         
         with gr.Row(equal_height=True):
             with gr.Column():
-                input_image = gr.Image(label="Upload Image", type="pil", height=290, sources=["upload", "webcam", "clipboard"])
-                prompt = gr.Text(label="Guiding Prompt", show_label=True, placeholder="e.g., a man with a red superhero mask")
+                input_image = gr.Image(
+                    label="Upload Image", 
+                    type="pil", 
+                    height=290, 
+                    sources=["upload", "webcam", "clipboard"]
+                )
+                prompt = gr.Text(
+                    label="Guiding Prompt", 
+                    show_label=True, 
+                    placeholder="e.g., a man with a red superhero mask"
+                )
                 
                 lora_adapter = gr.Dropdown(
                     label="Choose a Creative Style",
-                    # Updated choices for the new adapters
-                    choices=["Simple-Tuner", "Klein-Delight-Style"],
+                    choices=list(ADAPTER_SPECS.keys()),
                     value="Klein-Delight-Style"
                 )
 
@@ -160,18 +224,17 @@ with gr.Blocks() as demo:
                 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)
-                    # Updated defaults suitable for FLUX.2 Klein
-                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=4.0)       
+                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0)       
                     steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
                     
             with gr.Column():
                 output_image = gr.Image(label="Stylized Image", interactive=False, format="png", height=450)
                 used_seed = gr.Textbox(label="Used Seed", interactive=False)
 
-        # Updated examples for the new LoRAs
         gr.Examples(
             examples=[
                 ["examples/animal.jpg", "a cute red panda, charming and delightful illustration, soft lighting", "Klein-Delight-Style"],
+                ["examples/portrait.jpg", "cinematic lighting, high contrast", "Simple-Tuner"],
             ],
             inputs=[input_image, prompt, lora_adapter],
             outputs=[output_image, used_seed],