Update app.py

app.py

@@ -1,270 +1,192 @@
 import os
-import gc
 import gradio as gr
 import numpy as np
-import random
 import spaces
 import torch
-from diffusers import Flux2KleinPipeline
+import random
 from PIL import Image
+from typing import Iterable
 
-# --- Setup ---
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# 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
+from gradio.themes.utils import colors, fonts, sizes
+
+colors.orange_red = colors.Color(
+    name="orange_red", c50="#FFF0E5", c100="#FFE0CC", c200="#FFC299", c300="#FFA366",
+    c400="#FF8533", c500="#FF4500", c600="#E63E00", c700="#CC3700", c800="#B33000",
+    c900="#992900", c950="#802200",
+)
+
+class OrangeRedTheme(Soft):
+    def __init__(
+        self, *, primary_hue: colors.Color | str = colors.gray,
+        secondary_hue: colors.Color | str = colors.orange_red,
+        neutral_hue: colors.Color | str = colors.slate, text_size: sizes.Size | str = sizes.text_lg,
+        font: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
+        ),
+        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
+        ),
+    ):
+        super().__init__(
+            primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue,
+            text_size=text_size, font=font, font_mono=font_mono,
+        )
+        super().set(
+            background_fill_primary="*primary_50",
+            background_fill_primary_dark="*primary_900",
+            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
+            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
+            button_primary_text_color="white",
+            button_primary_text_color_hover="white",
+            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
+            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
+            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            slider_color="*secondary_500",
+            slider_color_dark="*secondary_600",
+            block_title_text_weight="600", block_border_width="3px",
+            block_shadow="*shadow_drop_lg", button_primary_shadow="*shadow_drop_lg",
+            button_large_padding="11px", color_accent_soft="*primary_100",
+            block_label_background_fill="*primary_200",
+        )
 
+orange_red_theme = OrangeRedTheme()
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
 
 # --- Model Loading ---
-# Only load the 9B distilled model as requested.
-REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-9B"
-
-print(f"Loading FLUX.2 Klein Distilled Model: {REPO_ID_DISTILLED}...")
-pipe = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype)
-pipe.to(device)
+print("Loading FLUX.2 Klein 9B model...")
+pipe = Flux2KleinPipeline.from_pretrained(
+    "black-forest-labs/FLUX.2-klein-9B",
+    torch_dtype=torch.bfloat16
+).to(device)
 print("Model loaded successfully.")
 
-# --- LoRA Adapter Configuration ---
-ADAPTER_SPECS = {
-    "Arcane": {
-        "repo": "DeverStyle/Flux.2-Klein-Loras",
-        "weights": "dever_arcane_flux2_klein_9b.safetensors",
-        "adapter_name": "Arcane"
-    },  
-     "Zoom": {
-        "repo": "fal/flux-2-klein-4B-zoom-lora",
-        "weights": "flux-red-zoom-lora.safetensors",
-        "adapter_name": "zoom"
-    }, 
-    "Background-Remove": {
-        "repo": "fal/flux-2-klein-4B-background-remove-lora",
-        "weights": "flux-background-remove-lora.safetensors",
-        "adapter_name": "rmbg"
-    },
-     "Object-Remove": {
-        "repo": "fal/flux-2-klein-4B-object-remove-lora",
-        "weights": "flux-object-remove-lora.safetensors",
-        "adapter_name": "object-remove"
-    }, 
-    "Sprite-Sheet": {
-        "repo": "fal/flux-2-klein-4b-spritesheet-lora",
-        "weights": "flux-spritesheet-lora.safetensors",
-        "adapter_name": "spritesheet"
-    },
+# --- LoRA Loading (Updated) ---
+print("Loading new LoRA adapters...")
+pipe.load_lora_weights(
+    "starsfriday/FLUX.2-klein-AC-Style-LORA",
+    weight_name="flux2_klein_lowres.safetensors",
+    adapter_name="american_comic_style"
+)
+pipe.load_lora_weights(
+    "linoyts/Flux2-Klein-Delight-LoRA",
+    weight_name="pytorch_lora_weights_v2.safetensors",
+    adapter_name="klein-delight"
+)
+print("All LoRA adapters loaded.")
+
+# Updated map for the new adapters
+ADAPTER_MAP = {
+    "American Comic Style": "american_comic_style",
+    "Klein Delight Style": "klein-delight",
 }
 
-# Keep track of which adapters have been downloaded and loaded
-LOADED_ADAPTERS = set()
-
-# --- Helper Functions ---
-def update_dimensions_from_image(image_list):
-    """
-    Update width/height sliders based on the first uploaded image's aspect ratio.
-    """
-    if not image_list:
-        return 1024, 1024
-    
-    # Gallery returns a list of tuples: [(<PIL.Image.Image>,), ...]
-    img = image_list[0][0] 
-    img_width, img_height = img.size
-    
-    aspect_ratio = img_width / img_height
-    
-    if aspect_ratio >= 1:
-        new_width = MAX_IMAGE_SIZE
-        new_height = int(MAX_IMAGE_SIZE / aspect_ratio)
-    else:
-        new_height = MAX_IMAGE_SIZE
-        new_width = int(MAX_IMAGE_SIZE * aspect_ratio)
-    
-    # Ensure dimensions are multiples of 8
-    new_width = (new_width // 8) * 8
-    new_height = (new_height // 8) * 8
-    
-    return new_width, new_height
-
-# --- Core Inference Function ---
-@spaces.GPU(duration=90)
-def infer(
-    prompt: str,
-    lora_adapter: str,
-    input_images=None,
-    seed: int = 42,
-    randomize_seed: bool = True,
-    width: int = 1024,
-    height: int = 1024,
-    num_inference_steps: int = 4,
-    guidance_scale: float = 1.0,
-    progress=gr.Progress(track_tqdm=True)
-):
-    """
-    Main function to generate or edit images using the FLUX.2 model and selected LoRA adapters.
-    """
-    gc.collect()
-    torch.cuda.empty_cache()
-
-    # --- LoRA Handling ---
-    if lora_adapter != "None":
-        spec = ADAPTER_SPECS.get(lora_adapter)
-        if not spec:
-            raise gr.Error(f"Configuration not found for adapter: {lora_adapter}")
-
-        adapter_name = spec["adapter_name"]
+@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
+    if not input_image:
+        raise gr.Error("Please upload an image to apply a style to.")
 
-        # Download and load the adapter if it's the first time being used
-        if adapter_name not in LOADED_ADAPTERS:
-            print(f"--- Downloading and Loading Adapter: {lora_adapter} ---")
-            progress(0.1, desc=f"Loading LoRA: {lora_adapter}")
-            try:
-                pipe.load_lora_weights(
-                    spec["repo"], 
-                    weight_name=spec["weights"], 
-                    adapter_name=adapter_name
-                )
-                LOADED_ADAPTERS.add(adapter_name)
-                print(f"--- Adapter {lora_adapter} loaded successfully. ---")
-            except Exception as e:
-                raise gr.Error(f"Failed to load adapter {lora_adapter}: {e}")
-        
-        # Set the active adapter
+    # 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})")
         pipe.set_adapters([adapter_name], adapter_weights=[1.0])
-
     else:
-        # If "None" is selected, ensure no LoRAs are active
+        # If "None" is selected (or an invalid choice), disable LoRAs
+        print("No LoRA selected. Disabling adapters.")
         pipe.disable_lora()
-
-    # --- Seed ---
+        
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-
-    # --- Image Processing ---
-    # Prepare image list from Gradio Gallery input
-    image_list = None
-    if input_images:
-        image_list = [item[0] for item in input_images] # Extract PIL images from tuples
-
-    # --- Generation ---
-    progress(0.5, desc=f"Generating with seed {seed}...")
     
-    pipe_kwargs = {
-        "prompt": prompt,
-        "height": height,
-        "width": width,
-        "num_inference_steps": num_inference_steps,
-        "guidance_scale": guidance_scale,
-        "generator": generator,
-    }
+    original_image = input_image.copy().convert("RGB")
     
-    # Add images to the pipeline arguments only if they are provided
-    if image_list:
-        pipe_kwargs["image"] = image_list
+    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]
 
-    try:
-        image = pipe(**pipe_kwargs).images[0]
-    except Exception as e:
-        # Unload the active LoRA on failure to allow retries or switching
-        pipe.disable_lora()
-        raise gr.Error(f"Inference failed: {e}")
-    
-    # --- Cleanup ---
-    gc.collect()
-    torch.cuda.empty_cache()
-    
+    return image, seed
+
+@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)
     return image, seed
 
 # --- UI Layout ---
-css = """
-#col-container { margin: 0 auto; max-width: 1200px; }
-.gallery-container img { object-fit: contain; }
+css="""
+#col-container { margin: 0 auto; max-width: 960px; }
+#main-title h1 { font-size: 2.2em !important; }
 """
 
 with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# **FLUX.2-klein-LoRA-Studio**")
+        gr.Markdown("# **FLUX.2 Klein LoRA Stylizer**", elem_id="main-title")
         gr.Markdown(
-            "Generate and edit images using the distilled **FLUX.2-klein-9B** model. "
-            "Select a specialized LoRA adapter from the dropdown for advanced editing tasks."
+            "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."
         )
         
-        with gr.Row():
-            with gr.Column(scale=2):
-                with gr.Row():
-                    prompt = gr.Text(
-                        label="Prompt",
-                        show_label=False,
-                        max_lines=2,
-                        placeholder="Enter your prompt here...",
-                        container=False,
-                        scale=3
-                    )
-                    run_button = gr.Button("Run", scale=1, variant="primary")
-                    
-                with gr.Accordion("Input Image(s) (for editing LoRAs)", open=False):
-                    input_images = gr.Gallery(
-                        label="Input Image(s)",
-                        type="pil",
-                        columns=3,
-                        rows=1,
-                    )
+        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")
                 
                 lora_adapter = gr.Dropdown(
-                    label="Select LoRA Adapter (or None for base generation)",
-                    choices=["None"] + list(ADAPTER_SPECS.keys()),
-                    value="None"
+                    label="Choose a Creative Style",
+                    # Updated choices for the new adapters
+                    choices=["American Comic Style", "Klein Delight Style"],
+                    value="American Comic Style"
                 )
-                
+
+                run_button = gr.Button("Apply Style", variant="primary")
+
                 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)       
+                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
                     
-                    with gr.Row():
-                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=8, value=1024)
-                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=8, value=1024)
-                    
-                    with gr.Row():
-                        num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=4)
-                        guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0)
-            
-            with gr.Column(scale=1):
-                result = gr.Image(label="Result", show_label=False, height=512)
+            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=[
-                ["A photorealistic, soaking wet capybara taking shelter under a large banana leaf in a rainy jungle, close up photo", "None"],
-                ["A kawaii die-cut sticker of a chubby orange cat, big sparkly eyes, happy smile, paws raised, heart-shaped pink nose, smooth rounded lines, black outlines, soft gradient shading, pink cheeks.", "None"],
-                ["A beautiful, majestic white horse running on a beach, cinematic lighting", "Zoom"],
-                ["A corgi wearing a tiny backpack, hiking in the mountains", "Background-Remove"],
-                ["A cute, round, fluffy creature with big eyes", "Sprite-Sheet"],
+                ["examples/portrait_man.jpg", "a man with a rugged beard, pop art style, bold lines, heavy shading", "American Comic Style"],
+                ["examples/cityscape.jpg", "a futuristic city, vibrant colors, clean lines, delightful style", "Klein Delight Style"],
+                ["examples/portrait_woman.jpg", "a woman with glasses, comic book art, detailed ink work, speech bubble", "American Comic Style"],
+                ["examples/animal.jpg", "a cute red panda, charming and delightful illustration, soft lighting", "Klein Delight Style"],
             ],
-            fn=infer,
-            inputs=[prompt, lora_adapter],
-            outputs=[result, used_seed],
+            inputs=[input_image, prompt, lora_adapter],
+            outputs=[output_image, used_seed],
+            fn=infer_example,
             cache_examples=False,
-            label="Examples"
         )
-
-    # --- Event Listeners ---
-    # Trigger inference on button click or prompt submission
+            
     run_button.click(
         fn=infer,
-        inputs=[prompt, lora_adapter, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale],
-        outputs=[result, used_seed]
-    )
-    prompt.submit(
-        fn=infer,
-        inputs=[prompt, lora_adapter, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale],
-        outputs=[result, used_seed]
-    )
-    
-    # Auto-update dimensions when an image is uploaded for editing
-    input_images.upload(
-        fn=update_dimensions_from_image,
-        inputs=[input_images],
-        outputs=[width, height]
+        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
+        outputs=[output_image, used_seed]
     )
 
-# --- Launch the App ---
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(css=css)
+    demo.queue().launch(css=css, theme=orange_red_theme, show_error=True)