Upload 10 files

README.md

@@ -1,13 +1,12 @@
 ---
-title: Qwen Image Edit 2511
-emoji: 🏆
-colorFrom: pink
-colorTo: red
+title: FLUX.2 [Klein] 4B
+emoji: 💻
+colorFrom: blue
+colorTo: gray
 sdk: gradio
-sdk_version: 6.2.0
+sdk_version: 5.29.1
 app_file: app.py
-pinned: false
-license: apache-2.0
+pinned: true
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,464 +1,465 @@
-import gradio as gr
-import numpy as np
-import random
-import math
-import time
-import torch
-import spaces
-
-from datetime import datetime
-import tempfile
-import zipfile
-from pathlib import Path
-from PIL import Image
-from diffusers import QwenImageEditPlusPipeline
-
-import os
-import base64
-import json
-
-SYSTEM_PROMPT = '''
-# Edit Prompt Enhancer
-You are a professional edit prompt enhancer. Your task is to generate a direct and specific edit prompt based on the user-provided instruction and the image input conditions.  
-
-Please strictly follow the enhancing rules below:
-
-## 1. General Principles
-- Keep the enhanced prompt **direct and specific**.  
-- If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary.  
-- Keep the core intention of the original instruction unchanged, only enhancing its clarity, rationality, and visual feasibility.  
-- All added objects or modifications must align with the logic and style of the edited input image’s overall scene.  
-
-## 2. Task-Type Handling Rules
-### 1. Add, Delete, Replace Tasks
-- If the instruction is clear (already includes task type, target entity, position, quantity, attributes), preserve the original intent and only refine the grammar.  
-- If the description is vague, supplement with minimal but sufficient details (category, color, size, orientation, position, etc.). For example:  
-    > Original: "Add an animal"  
-    > Rewritten: "Add a light-gray cat in the bottom-right corner, sitting and facing the camera"  
-- Remove meaningless instructions: e.g., "Add 0 objects" should be ignored or flagged as invalid.  
-- For replacement tasks, specify "Replace Y with X" and briefly describe the key visual features of X.  
-
-### 2. Text Editing Tasks
-- All text content must be enclosed in English double quotes `" "`. Keep the original language of the text, and keep the capitalization.  
-- Both adding new text and replacing existing text are text replacement tasks, For example:  
-    - Replace "xx" to "yy"  
-    - Replace the mask / bounding box to "yy"  
-    - Replace the visual object to "yy"  
-- Specify text position, color, and layout only if user has required.  
-- If font is specified, keep the original language of the font.  
-
-### 3. Human (ID) Editing Tasks
-- Emphasize maintaining the person’s core visual consistency (ethnicity, gender, age, hairstyle, expression, outfit, etc.).  
-- If modifying appearance (e.g., clothes, hairstyle), ensure the new element is consistent with the original style.  
-- **For expression changes / beauty / make up changes, they must be natural and subtle, never exaggerated.**  
-- Example:  
-    > Original: "Change the person’s hat"  
-    > Rewritten: "Replace the man’s hat with a dark brown beret; keep smile, short hair, and gray jacket unchanged"  
-
-### 4. Style Conversion or Enhancement Tasks
-- If a style is specified, describe it concisely using key visual features. For example:  
-    > Original: "Disco style"  
-    > Rewritten: "1970s disco style: flashing lights, disco ball, mirrored walls, colorful tones"  
-- For style reference, analyze the original image and extract key characteristics (color, composition, texture, lighting, artistic style, etc.), integrating them into the instruction.  
-- **Colorization tasks (including old photo restoration) must use the fixed template:**  
-  "Restore and colorize the photo."  
-- Clearly specify the object to be modified. For example:  
-    > Original: Modify the subject in Picture 1 to match the style of Picture 2.  
-    > Rewritten: Change the girl in Picture 1 to the ink-wash style of Picture 2 — rendered in black-and-white watercolor with soft color transitions.
-
-- If there are other changes, place the style description at the end.
-
-### 5. Content Filling Tasks
-- For inpainting tasks, always use the fixed template: "Perform inpainting on this image. The original caption is: ".
-- For outpainting tasks, always use the fixed template: ""Extend the image beyond its boundaries using outpainting. The original caption is: ".
-
-### 6. Multi-Image Tasks
-- Rewritten prompts must clearly point out which image’s element is being modified. For example:  
-    > Original: "Replace the subject of picture 1 with the subject of picture 2"  
-    > Rewritten: "Replace the girl of picture 1 with the boy of picture 2, keeping picture 2’s background unchanged"  
-- For stylization tasks, describe the reference image’s style in the rewritten prompt, while preserving the visual content of the source image.  
-
-## 3. Rationale and Logic Checks
-- Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" should be logically corrected.  
-- Add missing key information: e.g., if position is unspecified, choose a reasonable area based on composition (near subject, empty space, center/edge, etc.).  
-
-# Output Format Example
-```json
-{
-   "Rewritten": "..."
-}
-'''
-
-DEFAULT_TRUE_GUIDANCE_SCALE = 4.0
-DEFAULT_NUM_INFERENCE_STEPS = 40
-
-prompt_debug_value = [None]
-input_images_debug_value = [None]
-number_debug_value = [None]
-
-
-def encode_image(pil_image):
-    import io
-    buffered = io.BytesIO()
-    pil_image.save(buffered, format="PNG")
-    return base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-# --- Model Loading ---
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load the model pipeline
-pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2511", torch_dtype=dtype).to(device)
-
-# --- UI Constants and Helpers ---
-MAX_SEED = np.iinfo(np.int32).max
-
-# --- Main Inference Function (with hardcoded negative prompt) ---
-def infer(
-    images,
-    prompt,
-    seed=42,
-    randomize_seed=True,
-    true_guidance_scale=DEFAULT_TRUE_GUIDANCE_SCALE,
-    num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS,
-    height=None,
-    width=None,
-    rewrite_prompt=False,
-    num_images_per_prompt=1,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """
-    Generates an image using the local Qwen-Image diffusers pipeline.
-    """
-    # Hardcode the negative prompt as requested
-    negative_prompt = " "
-    
-    # Load input images into PIL Images
-    pil_images = []
-    if images is not None:
-        for item in images:
-            try:
-                if isinstance(item[0], Image.Image):
-                    pil_images.append(item[0].convert("RGB"))
-                elif isinstance(item[0], str):
-                    pil_images.append(Image.open(item[0]).convert("RGB"))
-                elif hasattr(item, "name"):
-                    pil_images.append(Image.open(item.name).convert("RGB"))
-            except Exception:
-                continue
-
-    if height==256 and width==256:
-        height, width = None, None
-    
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    print(f"Calling pipeline with prompt: '{prompt}'")
-    print(f"Negative Prompt: '{negative_prompt}'")
-    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
-    
-    return infer_on_gpu(
-        pil_images if len(pil_images) > 0 else None,
-        prompt,
-        negative_prompt,
-        seed,
-        randomize_seed,
-        true_guidance_scale,
-        num_inference_steps,
-        height,
-        width,
-        rewrite_prompt,
-        num_images_per_prompt,
-        progress
-    ), seed
-
-def get_duration(
-    pil_images,
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    true_guidance_scale,
-    num_inference_steps,
-    height,
-    width,
-    rewrite_prompt,
-    num_images_per_prompt,
-    progress,
-):
-    return 180 + (len(pil_images if pil_images is not None else 0) * 60)
-
-@spaces.GPU(duration=get_duration)
-def infer_on_gpu(
-    pil_images,
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    true_guidance_scale,
-    num_inference_steps,
-    height,
-    width,
-    rewrite_prompt,
-    num_images_per_prompt,
-    progress,
-):
-    # Set up the generator for reproducibility
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    # Generate the image
-    output_images = pipe(
-        image=pil_images,
-        prompt=prompt,
-        height=height,
-        width=width,
-        negative_prompt=negative_prompt,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-        true_cfg_scale=true_guidance_scale,
-        num_images_per_prompt=num_images_per_prompt,
-    ).images
-
-    return output_images
-
-def export_images_to_zip(gallery) -> str:
-    """
-    Bundle compiled_transformer_1 and compiled_transformer_2 into a zip file and return the file path.
-    """
-
-    tmp_zip = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
-    tmp_zip.close()
-
-    with zipfile.ZipFile(tmp_zip.name, "w", compression=zipfile.ZIP_DEFLATED) as zf:
-        for i in range(len(gallery)):
-            image_path = gallery[i]
-            zf.write(image_path, arcname=os.path.basename(image_path))
-
-    print(str(len(gallery)) + " images zipped")
-    return tmp_zip.name
-    
-def save_on_path(img: Image, filename: str, format_: str = None) -> Path:
-    tmp_dir = Path(tempfile.mkdtemp(prefix="pil_tmp_"))
-    file_path = tmp_dir / filename
-    if isinstance(img, np.ndarray):
-        img = Image.fromarray(img)
-    img.save(file_path, format=format_ or img.format)
-    
-    return file_path
-
-def infer_example(
-    input_images,
-    prompt,
-    seed=42,
-    randomize_seed=True,
-    true_guidance_scale=DEFAULT_TRUE_GUIDANCE_SCALE,
-    num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS,
-    height=None,
-    width=None,
-    rewrite_prompt=False,
-    num_images_per_prompt=1,
-):
-    start = time.time()
-    number=1
-    if prompt_debug_value[0] is not None or input_images_debug_value[0] is not None or number_debug_value[0] is not None:
-        prompt=prompt_debug_value[0]
-        input_images=input_images_debug_value[0]
-        number=number_debug_value[0]
-
-    gallery = []
-    for i in range(number):
-        try:
-            print("Generating #" + str(i + 1) + " image...")
-            seed = random.randint(0, MAX_SEED)
-            [output_images, seed] = infer(
-                input_images,
-                prompt,
-                seed,
-                randomize_seed,
-                true_guidance_scale,
-                num_inference_steps,
-                height,
-                width,
-                rewrite_prompt,
-                num_images_per_prompt
-            )
-            for output_image in output_images:
-                image_filename = datetime.now().strftime("%Y-%m-%d_%H-%M-%S.%f") + '.webp'
-                path = save_on_path(output_image, image_filename, format_="WEBP")
-                print("Image #" + str(i + 1) + " generated: " + str(path))
-                gallery.append(path)
-        except Exception as e:
-            print('Error: ' + e.message if e is not None else '')
-            #raise e
-    zip_path = export_images_to_zip(gallery)
-    print("ZIP path: " + str(zip_path))
-    
-    end = time.time()
-    secondes = int(end - start)
-    minutes = math.floor(secondes / 60)
-    secondes = secondes - (minutes * 60)
-    hours = math.floor(minutes / 60)
-    minutes = minutes - (hours * 60)
-    information = ("Start the process again if you want a different result. " if randomize_seed else "") + \
-    "The images have been generated in " + \
-    ((str(hours) + " h, ") if hours != 0 else "") + \
-    ((str(minutes) + " min, ") if hours != 0 or minutes != 0 else "") + \
-    str(secondes) + " sec (including " + str("allocation_time") + " seconds of GPU). " + \
-    "The images have " + str(num_images_per_prompt) + " num_images_per_prompt. " + \
-    "The image resolution is " + str(width) + \
-    " pixels large and " + str(height) + \
-    " pixels high" + ((", so a resolution of " + f'{width * height:,}' + " pixels") if width is not None and height is not None else "") + "."
-    return [seed, gallery, zip_path, information]
-
-# --- Examples and UI Layout ---
-examples = [
-    [["kill_bill.jpeg"], "A brunette woman"]
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 1024px;
-}
-#edit_text{margin-top: -62px !important}
-#default_examples {
-display:none;
-}
-"""
-
-js = """
-function afterGeneration() {
-  document.getElementById('download_btn').click();
-  return 0;
-}
-"""
-
-with gr.Blocks(css=css, js=js) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.HTML('<img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/qwen_image_edit_logo.png" alt="Qwen-Image Logo" width="400" style="display: block; margin: 0 auto;">')
-        gr.Markdown("[Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series. Try on [Qwen Chat](https://chat.qwen.ai/), or [download model](https://huggingface.co/Qwen/Qwen-Image-Edit) to run locally with ComfyUI or diffusers.")
-        with gr.Row():
-            with gr.Column():
-                input_images = gr.Gallery(label="Input Images", show_label=False, type="pil", interactive=True)
-
-            # result = gr.Image(label="Result", show_label=False, type="pil")
-            result = gr.Gallery(label="Result", show_label=False, type="pil")
-        with gr.Row():
-            prompt = gr.Text(
-                    label="Prompt",
-                    show_label=False,
-                    placeholder="describe the edit instruction",
-                    container=False,
-            )
-            run_button = gr.Button("🚀 Edit!", variant="primary")
-
-        with gr.Accordion("Advanced Settings", open=False):
-            # Negative prompt UI element is removed here
-
-            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():
-
-                true_guidance_scale = gr.Slider(
-                    label="True guidance scale",
-                    minimum=1.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=DEFAULT_TRUE_GUIDANCE_SCALE
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=60,
-                    step=1,
-                    value=DEFAULT_NUM_INFERENCE_STEPS
-                )
-                
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=2048,
-                    step=8,
-                    value=1024,
-                )
-                
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=2048,
-                    step=8,
-                    value=1024,
-                )
-                
-                
-                rewrite_prompt = gr.Checkbox(label="Rewrite prompt", value=True)
-
-    with gr.Row(elem_id="default_examples"):
-        prompt_debug = gr.Text(
-            max_lines=2,
-            container=False,
-            scale=3
-        )
-        input_images_debug = gr.Gallery(
-            label="Input Image(s)",
-            type="pil",
-            columns=3,
-            rows=1,
-        )
-        number_debug=gr.Slider(minimum=1, maximum=50, step=1, value=50)
-        download_button = gr.DownloadButton(elem_id="download_btn", interactive = True)
-        info_debug = gr.HTML(value = "")
-        gr.Examples(
-            examples=examples,
-            inputs=[input_images, prompt],
-            outputs=[seed, result, download_button, info_debug],
-            fn=infer_example,
-            cache_examples=True
-        )
-
-    def handle_field_debug_change(prompt_debug_data, input_images_debug_data, number_debug_data):
-        prompt_debug_value[0] = prompt_debug_data
-        input_images_debug_value[0] = input_images_debug_data
-        number_debug_value[0] = number_debug_data
-        return []
-
-    inputs_debug=[prompt_debug, input_images_debug, number_debug]
-
-    prompt_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-    input_images_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-    number_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            input_images,
-            prompt,
-            seed,
-            randomize_seed,
-            true_guidance_scale,
-            num_inference_steps,
-            height,
-            width,
-            rewrite_prompt,
-        ],
-        outputs=[result, seed],
-    )
-    result.change(
-        fn=lambda text : [],
-        inputs=[prompt],
-        outputs=[],
-        js="afterGeneration()"
-    )
-
-if __name__ == "__main__":
-    demo.launch(mcp_server=True, share=True)
+import os
+from datetime import datetime
+import tempfile
+import zipfile
+from pathlib import Path
+import subprocess
+import sys
+import io
+import gradio as gr
+import numpy as np
+import random
+import spaces
+import torch
+from diffusers import Flux2KleinPipeline
+import requests
+from PIL import Image
+import json
+import base64
+from huggingface_hub import InferenceClient
+
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+hf_client = InferenceClient(
+    api_key=os.environ.get("HF_TOKEN"),
+)
+VLM_MODEL = "baidu/ERNIE-4.5-VL-424B-A47B-Base-PT"
+
+SYSTEM_PROMPT_TEXT_ONLY = """You are an expert prompt engineer for FLUX.2 by Black Forest Labs. Rewrite user prompts to be more descriptive while strictly preserving their core subject and intent.
+
+Guidelines:
+1. Structure: Keep structured inputs structured (enhance within fields). Convert natural language to detailed paragraphs.
+2. Details: Add concrete visual specifics - form, scale, textures, materials, lighting (quality, direction, color), shadows, spatial relationships, and environmental context.
+3. Text in Images: Put ALL text in quotation marks, matching the prompt's language. Always provide explicit quoted text for objects that would contain text in reality (signs, labels, screens, etc.) - without it, the model generates gibberish.
+
+Output only the revised prompt and nothing else."""
+
+SYSTEM_PROMPT_WITH_IMAGES = """You are FLUX.2 by Black Forest Labs, an image-editing expert. You convert editing requests into one concise instruction (50-80 words, ~30 for brief requests).
+
+Rules:
+- Single instruction only, no commentary
+- Use clear, analytical language (avoid "whimsical," "cascading," etc.)
+- Specify what changes AND what stays the same (face, lighting, composition)
+- Reference actual image elements
+- Turn negatives into positives ("don't change X" → "keep X")
+- Make abstractions concrete ("futuristic" → "glowing cyan neon, metallic panels")
+- Keep content PG-13
+
+Output only the final instruction in plain text and nothing else."""
+
+# Model repository IDs for 4B
+REPO_ID_REGULAR = "black-forest-labs/FLUX.2-klein-base-4B"
+REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-4B"
+
+# Load both 4B models
+print("Loading 4B Regular model...")
+pipe_regular = Flux2KleinPipeline.from_pretrained(REPO_ID_REGULAR, torch_dtype=dtype)
+pipe_regular.to("cuda")
+
+print("Loading 4B Distilled model...")
+pipe_distilled = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype)
+pipe_distilled.to("cuda")
+
+# Dictionary for easy access
+pipes = {
+    "Distilled (4 steps)": pipe_distilled,
+    "Base (50 steps)": pipe_regular,
+}
+
+# Default steps for each mode
+DEFAULT_STEPS = {
+    "Distilled (4 steps)": 4,
+    "Base (50 steps)": 50,
+}
+
+# Default CFG for each mode
+DEFAULT_CFG = {
+    "Distilled (4 steps)": 1.0,
+    "Base (50 steps)": 4.0,
+}
+
+prompt_debug_value = [None]
+input_images_debug_value = [None]
+number_debug_value = [None]
+
+def image_to_data_uri(img):
+    buffered = io.BytesIO()
+    img.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return f"data:image/png;base64,{img_str}"
+
+
+def upsample_prompt_logic(prompt, image_list):
+    try:
+        if image_list and len(image_list) > 0:
+            # Image + Text Editing Mode
+            system_content = SYSTEM_PROMPT_WITH_IMAGES
+            
+            # Construct user message with text and images
+            user_content = [{"type": "text", "text": prompt}]
+            
+            for img in image_list:
+                data_uri = image_to_data_uri(img)
+                user_content.append({
+                    "type": "image_url",
+                    "image_url": {"url": data_uri}
+                })
+                
+            messages = [
+                {"role": "system", "content": system_content},
+                {"role": "user", "content": user_content}
+            ]
+        else:
+            # Text Only Mode
+            system_content = SYSTEM_PROMPT_TEXT_ONLY
+            messages = [
+                {"role": "system", "content": system_content},
+                {"role": "user", "content": prompt}
+            ]
+
+        completion = hf_client.chat.completions.create(
+            model=VLM_MODEL,
+            messages=messages,
+            max_tokens=1024
+        )
+        
+        return completion.choices[0].message.content
+    except Exception as e:
+        print(f"Upsampling failed: {e}")
+        return prompt
+
+
+def update_dimensions_from_image(image_list):
+    """Update width/height sliders based on uploaded image aspect ratio.
+    Keeps one side at 1024 and scales the other proportionally, with both sides as multiples of 8."""
+    if image_list is None or len(image_list) == 0:
+        return 1024, 1024  # Default dimensions
+    
+    # Get the first image to determine dimensions
+    img = image_list[0][0]  # Gallery returns list of tuples (image, caption)
+    img_width, img_height = img.size
+    
+    aspect_ratio = img_width / img_height
+    
+    if aspect_ratio >= 1:  # Landscape or square
+        new_width = 1024
+        new_height = int(1024 / aspect_ratio)
+    else:  # Portrait
+        new_height = 1024
+        new_width = int(1024 * aspect_ratio)
+    
+    # Round to nearest multiple of 8
+    new_width = round(new_width / 8) * 8
+    new_height = round(new_height / 8) * 8
+    
+    # Ensure within valid range (minimum 256, maximum 1024)
+    new_width = max(256, min(1024, new_width))
+    new_height = max(256, min(1024, new_height))
+    
+    return new_width, new_height
+
+
+def update_steps_from_mode(mode_choice):
+    """Update the number of inference steps based on the selected mode."""
+    return DEFAULT_STEPS[mode_choice], DEFAULT_CFG[mode_choice]
+
+
+@spaces.GPU(duration=85)
+def infer(prompt, input_images=None, mode_choice="Distilled (4 steps)", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidance_scale=4.0, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
+    
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    # Select the appropriate pipeline based on mode choice
+    pipe = pipes[mode_choice]
+    
+    # Prepare image list (convert None or empty gallery to None)
+    image_list = None
+    if input_images is not None and len(input_images) > 0:
+        image_list = []
+        for item in input_images:
+            image_list.append(item[0])
+
+    # 1. Upsampling (Network bound)
+    final_prompt = prompt
+    if prompt_upsampling:
+        progress(0.1, desc="Upsampling prompt...")
+        final_prompt = upsample_prompt_logic(prompt, image_list)
+        print(f"Original Prompt: {prompt}")
+        print(f"Upsampled Prompt: {final_prompt}")
+
+    # 2. Image Generation
+    progress(0.2, desc=f"Generating image with 4B {mode_choice}...")
+    
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    pipe_kwargs = {
+        "prompt": final_prompt,
+        "height": height,
+        "width": width,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "generator": generator,
+    }
+    
+    # Add images if provided
+    if image_list is not None:
+        pipe_kwargs["image"] = image_list
+    
+    image = pipe(**pipe_kwargs).images[0]
+    
+    return image, seed
+
+def export_images_to_zip(gallery) -> str:
+    """
+    Bundle compiled_transformer_1 and compiled_transformer_2 into a zip file and return the file path.
+    """
+
+    tmp_zip = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
+    tmp_zip.close()
+
+    with zipfile.ZipFile(tmp_zip.name, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+        for i in range(len(gallery)):
+            image_path = gallery[i]
+            zf.write(image_path, arcname=os.path.basename(image_path))
+
+    print(str(len(gallery)) + " images zipped")
+    return tmp_zip.name
+    
+def save_on_path(img: Image, filename: str, format_: str = None) -> Path:
+    """
+    Save `img` in a unique temporary folder under the given `filename`
+    and return its absolute path.
+    """
+    # 1) unique temporary folder
+    tmp_dir = Path(tempfile.mkdtemp(prefix="pil_tmp_"))
+    
+    # 2) full path of the future file
+    file_path = tmp_dir / filename
+    
+    # 3) save
+    img.save(file_path, format=format_ or img.format)
+    
+    return file_path
+
+def infer_example(prompt, input_images=None, mode_choice="Distilled (4 steps)", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidance_scale=4.0, prompt_upsampling=False):
+    number=1
+    if prompt_debug_value[0] is not None or input_images_debug_value[0] is not None or number_debug_value[0] is not None:
+        prompt=prompt_debug_value[0]
+        input_images=input_images_debug_value[0]
+        number=number_debug_value[0]
+
+    gallery = []
+    for i in range(number):
+        try:
+            print("Generating #" + str(i + 1) + " image...")
+            seed = random.randint(0, MAX_SEED)
+            [image, seed] = infer(prompt, input_images, mode_choice, seed, True, width, height, num_inference_steps, guidance_scale, prompt_upsampling)
+            image_filename = datetime.now().strftime("%Y-%m-%d_%H-%M-%S.%f") + '.webp'
+            path = save_on_path(image, image_filename, format_="WEBP")
+            gallery.append(path)
+        except Exception as e:
+            print('Error: ' + e.message if e is not None else '')
+            #raise e
+    zip_path = export_images_to_zip(gallery)
+    return [seed, zip_path, "Done!"]
+    
+
+examples = [
+    ["Create a vase on a table in living room, the color of the vase is a gradient of color, starting with #02eb3c color and finishing with #edfa3c. The flowers inside the vase have the color #ff0088"],
+    ["Photorealistic infographic showing the complete Berlin TV Tower (Fernsehturm) from ground base to antenna tip, full vertical view with entire structure visible including concrete shaft, metallic sphere, and antenna spire. Slight upward perspective angle looking up toward the iconic sphere, perfectly centered on clean white background. Left side labels with thin horizontal connector lines: the text '368m' in extra large bold dark grey numerals (#2D3748) positioned at exactly the antenna tip with 'TOTAL HEIGHT' in small caps below. The text '207m' in extra large bold with 'TELECAFÉ' in small caps below, with connector line touching the sphere precisely at the window level. Right side label with horizontal connector line touching the sphere's equator: the text '32m' in extra large bold dark grey numerals with 'SPHERE DIAMETER' in small caps below. Bottom section arranged in three balanced columns: Left - Large text '986' in extra bold dark grey with 'STEPS' in caps below. Center - 'BERLIN TV TOWER' in bold caps with 'FERNSEHTURM' in lighter weight below. Right - 'INAUGURATED' in bold caps with 'OCTOBER 3, 1969' below. All typography in modern sans-serif font (such as Inter or Helvetica), color #2D3748, clean minimal technical diagram style. Horizontal connector lines are thin, precise, and clearly visible, touching the tower structure at exact corresponding measurement points. Professional architectural elevation drawing aesthetic with dynamic low angle perspective creating sense of height and grandeur, poster-ready infographic design with perfect visual hierarchy."],
+    ["Soaking wet capybara taking shelter under a banana leaf in the rainy jungle, close up photo"],
+    ["A kawaii die-cut sticker of a chubby orange cat, featuring big sparkly eyes and a happy smile with paws raised in greeting and a heart-shaped pink nose. The design should have smooth rounded lines with black outlines and soft gradient shading with pink cheeks."],
+]
+
+examples_images = [
+    ["The person from image 1 is petting the cat from image 2, the bird from image 3 is next to them", ["woman1.webp", "cat_window.webp", "bird.webp"]]
+]
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1200px;
+}
+.gallery-container img{
+    object-fit: contain;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""# FLUX.2 [Klein] - 4B (Apache 2.0)
+FLUX.2 [klein] is a fast, unified image generation and editing model designed for fast inference [[model](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B)], [[blog](https://bfl.ai/blog/flux-2)]
+        """)
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=2,
+                        placeholder="Enter your prompt",
+                        container=False,
+                        scale=3
+                    )
+                    
+                    run_button = gr.Button("Run", scale=1)
+                    
+                with gr.Accordion("Input image(s) (optional)", open=False):
+                    input_images = gr.Gallery(
+                        label="Input Image(s)",
+                        type="pil",
+                        columns=3,
+                        rows=1,
+                    )
+                    input_image = gr.Image(label="Upload the image for editing", type="pil")
+
+                mode_choice = gr.Radio(
+                    label="Mode",
+                    choices=["Distilled (4 steps)", "Base (50 steps)"],
+                    value="Distilled (4 steps)",
+                )
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    
+                    prompt_upsampling = gr.Checkbox(
+                        label="Prompt Upsampling",
+                        value=False,
+                        info="Automatically enhance the prompt using a VLM"
+                    )
+        
+                    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=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="Number of inference steps",
+                            minimum=1,
+                            maximum=100,
+                            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():
+                result = gr.Image(label="Result", show_label=False)
+            
+        
+        gr.Examples(
+            examples=examples,
+            fn=infer,
+            inputs=[prompt],
+            outputs=[result, seed],
+            cache_examples=True,
+            cache_mode="lazy"
+        )
+
+        gr.Examples(
+            examples=examples_images,
+            fn=infer,
+            inputs=[prompt, input_images],
+            outputs=[result, seed],
+            cache_examples=True,
+            cache_mode="lazy"
+        )
+
+    # Auto-update dimensions when images are uploaded
+    input_images.upload(
+        fn=update_dimensions_from_image,
+        inputs=[input_images],
+        outputs=[width, height]
+    )
+    
+    # Auto-update steps when mode changes
+    mode_choice.change(
+        fn=update_steps_from_mode,
+        inputs=[mode_choice],
+        outputs=[num_inference_steps, guidance_scale]
+    )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[prompt, input_images, mode_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
+        outputs=[result, seed]
+    )
+    
+    with gr.Row(visible=False):
+        download_button = gr.DownloadButton(elem_id="download_btn", interactive = True)
+        info_debug = gr.HTML(value = "")
+        prompt_debug = gr.Text(
+            max_lines=2,
+            container=False,
+            scale=3
+        )
+        input_images_debug = gr.Gallery(
+            label="Input Image(s)",
+            type="pil",
+            columns=3,
+            rows=1,
+        )
+        gr.Examples(
+            examples=[
+                ["A dog", "woman1.webp"]
+            ],
+            fn=infer_example,
+            inputs=[prompt, input_image],
+            outputs=[seed, download_button, info_debug],
+            run_on_click=True,
+            cache_examples=True,
+            cache_mode='lazy'
+        )
+        number_debug=gr.Slider(minimum=1, maximum=50, step=1, value=50)
+
+    def handle_field_debug_change(prompt_debug_data, input_images_debug_data, number_debug_data):
+        prompt_debug_value[0] = prompt_debug_data
+        input_images_debug_value[0] = input_images_debug_data
+        number_debug_value[0] = number_debug_data
+        return []
+
+    inputs_debug=[prompt_debug, input_images_debug, number_debug]
+
+    prompt_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
+    input_images_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
+    number_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
+
+demo.launch()

diffusers.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f21e689d4807674000c1da8d3a084f7d521e961ee26ebd77e7a55a8efc3b95d
+size 5269929

optimization.py

@@ -4,37 +4,34 @@
 from typing import Any
 from typing import Callable
 from typing import ParamSpec
-
-import os
 import spaces
 import torch
-from torch.utils._pytree import tree_map_only
-from torchao.quantization import quantize_
-from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
-from torchao.quantization import Int8WeightOnlyConfig
-from huggingface_hub import hf_hub_download
-
-from optimization_utils import capture_component_call
-from optimization_utils import aoti_compile
-from optimization_utils import drain_module_parameters
-from optimization_utils import ZeroGPUCompiledModelFromDict  # NEW
-
+from spaces.zero.torch.aoti import ZeroGPUCompiledModel
+from spaces.zero.torch.aoti import ZeroGPUWeights
+from torch.utils._pytree import tree_map
 
 P = ParamSpec('P')
 
-# Expose compiled models so app.py can offer them for download
-COMPILED_TRANSFORMER_1 = None
-COMPILED_TRANSFORMER_2 = None
-
-LATENT_FRAMES_DIM = torch.export.Dim('num_latent_frames', min=8, max=81)
-LATENT_PATCHED_HEIGHT_DIM = torch.export.Dim('latent_patched_height', min=30, max=52)
-LATENT_PATCHED_WIDTH_DIM = torch.export.Dim('latent_patched_width', min=30, max=52)
+TRANSFORMER_IMAGE_DIM = torch.export.Dim('image_seq_length', min=4096, max=16384) # min: 0 images, max: 3 (1024x1024) images
 
 TRANSFORMER_DYNAMIC_SHAPES = {
-    'hidden_states': {
-        2: LATENT_FRAMES_DIM,
-        3: 2 * LATENT_PATCHED_HEIGHT_DIM,
-        4: 2 * LATENT_PATCHED_WIDTH_DIM,
+    'double': {
+        'hidden_states': {
+            1: TRANSFORMER_IMAGE_DIM,
+        },
+        'image_rotary_emb': (
+            {0: TRANSFORMER_IMAGE_DIM + 512},
+            {0: TRANSFORMER_IMAGE_DIM + 512},
+        ),
+    },
+    'single': {
+        'hidden_states': {
+            1: TRANSFORMER_IMAGE_DIM + 512,
+        },
+        'image_rotary_emb': (
+            {0: TRANSFORMER_IMAGE_DIM + 512},
+            {0: TRANSFORMER_IMAGE_DIM + 512},
+        ),
     },
 }
 
@@ -47,110 +44,34 @@ INDUCTOR_CONFIGS = {
     'triton.cudagraphs': True,
 }
 
-
-def load_compiled_transformers_from_hub(
-    repo_id: str,
-    filename_1: str = "compiled_transformer_1.pt",
-    filename_2: str = "compiled_transformer_2.pt",
-    device: str = "cuda",
-):
-    """
-    Loads the payload dicts (created via ZeroGPUCompiledModel.to_serializable_dict() and torch.save)
-    and rebuilds callable models that will move constants to CUDA on first call.
-    """
-    path_1 = hf_hub_download(repo_id=repo_id, filename=filename_1)
-    path_2 = hf_hub_download(repo_id=repo_id, filename=filename_2)
-
-    payload_1 = torch.load(path_1, map_location="cpu", weights_only=False)
-    payload_2 = torch.load(path_2, map_location="cpu", weights_only=False)
-
-    if not isinstance(payload_1, dict) or not isinstance(payload_2, dict):
-        raise TypeError("Precompiled files are not payload dicts. Please re-export them with to_serializable_dict().")
-
-    compiled_1 = ZeroGPUCompiledModelFromDict(payload_1, device=device)
-    compiled_2 = ZeroGPUCompiledModelFromDict(payload_2, device=device)
-    return compiled_1, compiled_2
-
-
-def _strtobool(v: str | None, default: bool = True) -> bool:
-    if v is None:
-        return default
-    return v.strip().lower() in ("1", "true", "yes", "y", "on")
-
-
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
-    global COMPILED_TRANSFORMER_1, COMPILED_TRANSFORMER_2
 
-    @spaces.GPU(duration=1500)
-    def compile_transformer():
-        pipeline.load_lora_weights(
-            "Kijai/WanVideo_comfy",
-            weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
-            adapter_name="lightx2v",
-        )
-        kwargs_lora = {"load_into_transformer_2": True}
-        pipeline.load_lora_weights(
-            "Kijai/WanVideo_comfy",
-            weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
-            adapter_name="lightx2v_2",
-            **kwargs_lora,
-        )
-        pipeline.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1.0, 1.0])
-        pipeline.fuse_lora(adapter_names=["lightx2v"], lora_scale=3.0, components=["transformer"])
-        pipeline.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1.0, components=["transformer_2"])
-        pipeline.unload_lora_weights()
+    blocks = {
+        'double': pipeline.transformer.transformer_blocks,
+        'single': pipeline.transformer.single_transformer_blocks,
+    }
 
-        with capture_component_call(pipeline, "transformer") as call:
+    @spaces.GPU(duration=1200)
+    def compile_block(blocks_kind: str):
+        block = blocks[blocks_kind][0]
+        with spaces.aoti_capture(block) as call:
             pipeline(*args, **kwargs)
 
-        dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda t: None, call.kwargs)
-        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
-
-        quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
-        quantize_(pipeline.transformer_2, Float8DynamicActivationFloat8WeightConfig())
-
-        exported_1 = torch.export.export(
-            mod=pipeline.transformer,
-            args=call.args,
-            kwargs=call.kwargs,
-            dynamic_shapes=dynamic_shapes,
-        )
-        exported_2 = torch.export.export(
-            mod=pipeline.transformer_2,
-            args=call.args,
-            kwargs=call.kwargs,
-            dynamic_shapes=dynamic_shapes,
-        )
+        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES[blocks_kind]
 
-        compiled_1 = aoti_compile(exported_1, INDUCTOR_CONFIGS)
-        compiled_2 = aoti_compile(exported_2, INDUCTOR_CONFIGS)
-        return compiled_1, compiled_2
-
-    # Quantize text encoder
-    quantize_(pipeline.text_encoder, Int8WeightOnlyConfig())
-
-    use_precompiled = False
-    precompiled_repo = os.getenv("WAN_PRECOMPILED_REPO", "Fabrice-TIERCELIN/Wan_2.2_compiled")
-
-    if use_precompiled:
-        try:
-            compiled_transformer_1, compiled_transformer_2 = load_compiled_transformers_from_hub(
-                repo_id=precompiled_repo,
-                device="cuda",
+        with torch.no_grad():      
+            exported = torch.export.export(
+                mod=block,
+                args=call.args,
+                kwargs=call.kwargs,
+                dynamic_shapes=dynamic_shapes,
             )
-        except Exception as e:
-            # fallback if payload format is wrong / outdated
-            print(f"[WARN] Failed to load precompiled artifacts ({e}). Falling back to GPU compilation.")
-            compiled_transformer_1, compiled_transformer_2 = compile_transformer()
-    else:
-        compiled_transformer_1, compiled_transformer_2 = compile_transformer()
-
-    # expose for downloads
-    COMPILED_TRANSFORMER_1 = compiled_transformer_1
-    COMPILED_TRANSFORMER_2 = compiled_transformer_2
 
-    pipeline.transformer.forward = compiled_transformer_1
-    drain_module_parameters(pipeline.transformer)
+        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS).archive_file
 
-    pipeline.transformer_2.forward = compiled_transformer_2
-    drain_module_parameters(pipeline.transformer_2)
+    for blocks_kind in ('double', 'single'):
+        archive_file = compile_block(blocks_kind)
+        for block in blocks[blocks_kind]:
+            weights = ZeroGPUWeights(block.state_dict())
+            block.forward = ZeroGPUCompiledModel(archive_file, weights)

requirements.txt

@@ -1,8 +1,8 @@
-git+https://github.com/huggingface/diffusers.git@973a077c6a4e7e7a7ea61a84bedd29ac24fb609a
+git+https://github.com/huggingface/diffusers.git
 transformers
 accelerate
 safetensors
-sentencepiece
-dashscope
+bitsandbytes
+torchao
 kernels
-torchvision
+spaces==0.43.0