update app

app.py

@@ -0,0 +1,351 @@
+import os
+import gc
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import random
+from PIL import Image
+from typing import Iterable
+from gradio.themes import Soft
+from gradio.themes.utils import colors, fonts, sizes
+from diffusers import Flux2KleinPipeline
+
+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)",
+            button_secondary_text_color="black",
+            button_secondary_text_color_hover="white",
+            button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
+            button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
+            button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
+            button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
+            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()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dtype = torch.bfloat16
+
+print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
+print("torch.__version__ =", torch.__version__)
+print("Using device:", device)
+
+REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-4B"
+
+print(f"Loading 4B Distilled model from {REPO_ID_DISTILLED}...")
+try:
+    pipe = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype)
+    pipe.to(device)
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise e
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+ADAPTER_SPECS = {
+    "None": {
+        "repo": None,
+        "weights": None,
+        "adapter_name": "none"
+    },
+    "Outpaint": {
+        "repo": "fal/flux-2-klein-4B-outpaint-lora",
+        "weights": "flux-outpaint-lora.safetensors",
+        "adapter_name": "outpaint"
+    },  
+     "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"
+    },
+}
+
+CURRENT_LORA_NAME = "None"
+
+def update_dimensions_from_image(image_list):
+    """Update width/height based on upload to maintain aspect ratio."""
+    if image_list is None or not image_list:
+        return 1024, 1024
+    
+    try:
+        if isinstance(image_list[0], tuple) or isinstance(image_list[0], list):
+            img = image_list[0][0]
+        else:
+            img = image_list[0]
+            
+        if isinstance(img, str):
+            img = Image.open(img)
+            
+        img_width, img_height = img.size
+        
+        if img_width > img_height:
+            new_width = 1024
+            new_height = int(1024 * (img_height / img_width))
+        else:
+            new_height = 1024
+            new_width = int(1024 * (img_width / img_height))
+        
+        # Round to nearest 8 (Flux requirement)
+        new_width = round(new_width / 16) * 16
+        new_height = round(new_height / 16) * 16
+        
+        return new_width, new_height
+    except Exception as e:
+        print(f"Error sizing image: {e}")
+        return 1024, 1024
+
+@spaces.GPU
+def infer(
+    prompt, 
+    input_images, 
+    lora_choice, 
+    seed, 
+    randomize_seed, 
+    width, 
+    height, 
+    num_inference_steps, 
+    guidance_scale, 
+    progress=gr.Progress(track_tqdm=True)
+):
+    global CURRENT_LORA_NAME
+    
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    if lora_choice != CURRENT_LORA_NAME:
+        print(f"--- Switching Adapter: {CURRENT_LORA_NAME} -> {lora_choice} ---")
+        
+        if CURRENT_LORA_NAME != "None":
+            try:
+                pipe.unload_lora_weights()
+                print("Unloaded previous weights.")
+            except Exception as e:
+                print(f"Warning unloading weights: {e}")
+        
+        if lora_choice != "None":
+            spec = ADAPTER_SPECS[lora_choice]
+            try:
+                print(f"Loading {lora_choice} weights from {spec['repo']}...")
+                pipe.load_lora_weights(
+                    spec["repo"], 
+                    weight_name=spec["weights"], 
+                    adapter_name=spec["adapter_name"]
+                )
+            except Exception as e:
+                raise gr.Error(f"Failed to load adapter {lora_choice}: {e}")
+        
+        CURRENT_LORA_NAME = lora_choice
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    pil_images = []
+    if input_images:
+        for item in input_images:
+            if isinstance(item, tuple) or isinstance(item, list):
+                img_obj = item[0]
+            else:
+                img_obj = item
+                
+            if isinstance(img_obj, str):
+                pil_images.append(Image.open(img_obj).convert("RGB"))
+            elif isinstance(img_obj, Image.Image):
+                pil_images.append(img_obj.convert("RGB"))
+    
+    pipe_kwargs = {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "generator": generator,
+    }
+
+    if pil_images and len(pil_images) > 0:
+
+        pipe_kwargs["image"] = pil_images[0]
+
+        pipe_kwargs["strength"] = 0.85
+
+    try:
+        print(f"Generating with Prompt: {prompt[:50]}... | LoRA: {CURRENT_LORA_NAME}")
+        result_image = pipe(**pipe_kwargs).images[0]
+        return result_image, seed
+    except Exception as e:
+        raise gr.Error(f"Generation failed: {str(e)}")
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache()
+
+@spaces.GPU
+def infer_example(prompt, images, lora_choice):
+    return infer(
+        prompt=prompt,
+        input_images=images,
+        lora_choice=lora_choice,
+        seed=42,
+        randomize_seed=False,
+        width=1024,
+        height=1024,
+        num_inference_steps=4, 
+        guidance_scale=1.0    
+    )
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1000px;
+}
+#main-title h1 {font-size: 2.3em !important;}
+"""
+
+with gr.Blocks() as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("# **FLUX.2-klein-LoRA-Studio**", elem_id="main-title")
+        gr.Markdown("Generate and Edit using the unified **FLUX.2-klein-4B** (Distilled) model with specialized LoRA adapters.")
+
+        with gr.Row(equal_height=True):
+            with gr.Column():
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=True,
+                    placeholder="Describe your image or edit...",
+                    lines=3
+                )
+                
+                with gr.Row():
+                    lora_choice = gr.Dropdown(
+                        label="Select Adapter",
+                        choices=list(ADAPTER_SPECS.keys()),
+                        value="None",
+                        info="Select 'None' for raw 4B Distilled generation."
+                    )
+                    run_button = gr.Button("Run / Edit", variant="primary", scale=1)
+
+                input_images = gr.Gallery(
+                    label="Input Image (Optional for T2I, Required for Edits)", 
+                    type="pil", 
+                    columns=2, 
+                    rows=1, 
+                    height=250,
+                    allow_preview=True
+                )
+
+            with gr.Column():
+                output_image = gr.Image(label="Result", interactive=False, format="png")
+                
+                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)
+                    
+                    with gr.Row():
+                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=16, value=1024)
+                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=16, value=1024)
+                    
+                    with gr.Row():
+                        guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0, info="Keep low (1.0) for Distilled model")
+                        num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=4, info="4-8 steps recommended for Distilled")
+
+        gr.Examples(
+            examples=[
+                ["A futuristic city with glowing neon lights, cinematic view.", None, "None"],
+                ["Zoom out to show the table.", ["examples/cup.jpg"], "Zoom"],
+                ["Remove the background.", ["examples/object.jpg"], "Background-Remove"],
+                ["Make a sprite sheet of this character.", ["examples/char.jpg"], "Sprite-Sheet"],
+            ],
+            inputs=[prompt, input_images, lora_choice],
+            outputs=[output_image, seed],
+            # fn=infer_example, # Disabled cache_examples for speed in snippet
+            label="Examples (Note: Upload images for edit examples)"
+        )
+        
+        gr.Markdown("Running on **FLUX.2-klein-4B** (Distilled). Adapters provided by Fal.ai.")
+
+    input_images.upload(
+        fn=update_dimensions_from_image,
+        inputs=[input_images],
+        outputs=[width, height]
+    )
+
+    run_button.click(
+        fn=infer,
+        inputs=[prompt, input_images, lora_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale],
+        outputs=[output_image, seed]
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(css=css, theme=orange_red_theme, ssr_mode=False, show_error=True)