Upload 10 files

.gitattributes

@@ -70,3 +70,4 @@ poli_tower.png filter=lfs diff=lfs merge=lfs -text
 squatting_sonic.png filter=lfs diff=lfs merge=lfs -text
 tower_takes_off.png filter=lfs diff=lfs merge=lfs -text
 ugly_sonic.jpeg filter=lfs diff=lfs merge=lfs -text
+bird.webp filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,14 +1,12 @@
 ---
-title: FLUX.1 Kontext
-emoji: ⚡
-colorFrom: green
+title: FLUX.2 [Klein] 4B
+emoji: 💻
+colorFrom: blue
 colorTo: gray
 sdk: gradio
-sdk_version: 5.29.1
+sdk_version: 6.3.0
 app_file: app.py
 pinned: true
-license: mit
-short_description: 'Kontext image editing on FLUX[dev] '
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,198 +1,275 @@
-# PyTorch 2.8 (temporary hack)
 import os
-os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
-
-# Actual demo code
-try:
-    import spaces
-except:
-    class spaces():
-        def GPU(*args, **kwargs):
-            def decorator(function):
-                return lambda *dummy_args, **dummy_kwargs: function(*dummy_args, **dummy_kwargs)
-            return decorator
-
+import subprocess
+import sys
+import io
 import gradio as gr
 import numpy as np
-import torch
 import random
-import os
-from datetime import datetime
-
+import spaces
+import torch
+from diffusers import Flux2KleinPipeline
+import requests
 from PIL import Image
-import tempfile
-import zipfile
-import shutil
-from pathlib import Path
+import json
+import base64
+from huggingface_hub import InferenceClient
 
-from diffusers import FluxKontextPipeline
-from diffusers.utils import load_image
-
-from optimization import optimize_pipeline_
+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,
+}
 
-pipe = FluxKontextPipeline.from_pretrained("yuvraj108c/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
-optimize_pipeline_(pipe, image=Image.new("RGB", (512, 512)), prompt='prompt')
-
-input_image_debug_value = [None]
-prompt_debug_value = [None]
-number_debug_value = [None]
-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_"))
+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
     
-    # 2) full path of the future file
-    file_path = tmp_dir / filename
+    # 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
     
-    # 3) save
-    img.save(file_path, format=format_ or img.format)
+    aspect_ratio = img_width / img_height
     
-    return file_path
-
-@spaces.GPU(duration=40)
-def infer(
-    input_image,
-    prompt,
-    seed = 42,
-    randomize_seed = False,
-    guidance_scale = 2.5,
-    steps = 28,
-    width = -1,
-    height = -1,
-    progress=gr.Progress(track_tqdm=True)
-):
-    """
-    Perform image editing using the FLUX.1 Kontext pipeline.
+    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)
     
-    This function takes an input image and a text prompt to generate a modified version
-    of the image based on the provided instructions. It uses the FLUX.1 Kontext model
-    for contextual image editing tasks.
+    # Round to nearest multiple of 8
+    new_width = round(new_width / 8) * 8
+    new_height = round(new_height / 8) * 8
     
-    Args:
-        input_image (PIL.Image.Image): The input image to be edited. Will be converted
-            to RGB format if not already in that format.
-        prompt (str): Text description of the desired edit to apply to the image.
-            Examples: "Remove glasses", "Add a hat", "Change background to beach".
-        seed (int, optional): Random seed for reproducible generation. Defaults to 42.
-            Must be between 0 and MAX_SEED (2^31 - 1).
-        randomize_seed (bool, optional): If True, generates a random seed instead of
-            using the provided seed value. Defaults to False.
-        guidance_scale (float, optional): Controls how closely the model follows the
-            prompt. Higher values mean stronger adherence to the prompt but may reduce
-            image quality. Range: 1.0-10.0. Defaults to 2.5.
-        steps (int, optional): Controls how many steps to run the diffusion model for.
-            Range: 1-30. Defaults to 28.
-        progress (gr.Progress, optional): Gradio progress tracker for monitoring
-            generation progress. Defaults to gr.Progress(track_tqdm=True).
+    # Ensure within valid range (minimum 256, maximum 1024)
+    new_width = max(256, min(1024, new_width))
+    new_height = max(256, min(1024, new_height))
     
-    Returns:
-        tuple: A 3-tuple containing:
-            - PIL.Image.Image: The generated/edited image
-            - int: The seed value used for generation (useful when randomize_seed=True)
-            - gr.update: Gradio update object to make the reuse button visible
+    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)):
     
-    Example:
-        >>> edited_image, used_seed, button_update = infer(
-        ...     input_image=my_image,
-        ...     prompt="Add sunglasses",
-        ...     seed=123,
-        ...     randomize_seed=False,
-        ...     guidance_scale=2.5
-        ... )
-    """
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-    if input_image:
-        input_image = input_image.convert("RGB")
-        image = pipe(
-            image=input_image, 
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            width = input_image.size[0] if width == -1 else width,
-            height = input_image.size[1] if height == -1 else height,
-            num_inference_steps=steps,
-            generator=torch.Generator().manual_seed(seed),
-        ).images[0]
-    else:
-        image = pipe(
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=steps,
-            generator=torch.Generator().manual_seed(seed),
-        ).images[0]
+    # Select the appropriate pipeline based on mode choice
+    pipe = pipes[mode_choice]
     
-    image_filename = datetime.now().strftime("%Y-%m-%d_%H-%M-%S.%f") + '.webp'
-    path = save_on_path(image, image_filename, format_="WEBP")
-    return path, gr.update(value=path, visible=True), seed, gr.update(visible=True)
-
-def infer_example(input_image, prompt):
-    number=1
-    if input_image_debug_value[0] is not None or prompt_debug_value[0] is not None or number_debug_value[0] is not None:
-        input_image=input_image_debug_value[0]
-        prompt=prompt_debug_value[0]
-        number=number_debug_value[0]
-        #input_image_debug_value[0]=prompt_debug_value[0]=prompt_debug_value[0]=None
-    gallery = []
-    try:
-        for i in range(number):
-            print("Generating #" + str(i + 1) + " image...")
-            seed = random.randint(0, MAX_SEED)
-            image, download_button, seed, _ = infer(input_image, prompt, seed, True)
-            gallery.append(image)
-    except:
-        print("Error")
-    zip_path = export_images_to_zip(gallery)
-    return gallery, seed, zip_path
-
-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
-
-css="""
+    # 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
+
+
+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: 960px;
+    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.1 Kontext [dev]
-Image editing and manipulation model guidance-distilled from FLUX.1 Kontext [pro], [[blog]](https://bfl.ai/announcements/flux-1-kontext-dev) [[model]](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev)
+        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():
-                input_image = gr.Image(label="Upload the image for editing", type="pil")
                 with gr.Row():
                     prompt = gr.Text(
                         label="Prompt",
                         show_label=False,
-                        max_lines=1,
-                        placeholder="Enter your prompt for editing (e.g., 'Remove glasses', 'Add a hat')",
+                        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,
                     )
-                    run_button = gr.Button(value="🚀 Edit", variant = "primary", scale=0)
+
+                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,
@@ -203,97 +280,84 @@ Image editing and manipulation model guidance-distilled from FLUX.1 Kontext [pro
                     
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     
-                    guidance_scale = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=1,
-                        maximum=10,
-                        step=0.1,
-                        value=2.5,
-                    )       
-                    
-                    steps = gr.Slider(
-                        label="Steps",
-                        minimum=1,
-                        maximum=30,
-                        value=30,
-                        step=1
-                    )      
-                    
-                    width = gr.Slider(
-                        label="Output width",
-                        info="-1 = original width",
-                        minimum=-1,
-                        maximum=1024,
-                        value=-1,
-                        step=1
-                    )   
-                    
-                    height = gr.Slider(
-                        label="Output height",
-                        info="-1 = original height",
-                        minimum=-1,
-                        maximum=1024,
-                        value=-1,
-                        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="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, interactive=False)
-                download_button = gr.DownloadButton(elem_id="download_btn", visible=False)
-                reuse_button = gr.Button("Reuse this image", visible=False)
-    
-    with gr.Row(visible=False):
-        download_button = gr.DownloadButton(elem_id="download_btn", interactive = True)
-        result_gallery = gr.Gallery(label = 'Downloadable results', show_label = True, interactive = False, elem_id = "gallery1")
+                result = gr.Image(label="Result", show_label=False)
+            
+        
         gr.Examples(
-            examples=[
-                ["monster.png", "Make this monster ride a skateboard on the beach"]
-            ],
-            inputs=[input_image, prompt],
-            outputs=[result_gallery, seed, download_button],
-            fn=infer_example,
-            run_on_click=True,
+            examples=examples,
+            fn=infer,
+            inputs=[prompt],
+            outputs=[result, seed],
             cache_examples=True,
-            cache_mode='lazy'
+            cache_mode="lazy"
         )
-        prompt_debug=gr.Textbox(label="Prompt Debug")
-        input_image_debug=gr.Image(type="pil", label="Image Debug")
-        number_debug=gr.Slider(label="Number Debug", minimum=1, maximum=50, step=1, value=50)
-
-    gr.Examples(
-        label = "Examples from demo",
-        examples=[
-            ["flowers.png", "turn the flowers into sunflowers"],
-            ["monster.png", "make this monster ride a skateboard on the beach"],
-            ["cat.png", "make this cat happy"]
-        ],
-        inputs=[input_image, prompt],
-        outputs=[result, download_button, seed],
-        fn=infer
-    )
 
-    def handle_field_debug_change(input_image_debug_data, prompt_debug_data, number_debug_data):
-        prompt_debug_value[0] = prompt_debug_data
-        input_image_debug_value[0] = input_image_debug_data
-        number_debug_value[0] = number_debug_data
-        return []
+        gr.Examples(
+            examples=examples_images,
+            fn=infer,
+            inputs=[prompt, input_images],
+            outputs=[result, seed],
+            cache_examples=True,
+            cache_mode="lazy"
+        )
 
-    inputs_debug=[input_image_debug, prompt_debug, number_debug]
+    # 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]
+    )
 
-    input_image_debug.upload(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-    prompt_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_image, prompt, seed, randomize_seed, guidance_scale, steps, width, height],
-        outputs = [result, download_button, seed, reuse_button]
-    )
-    reuse_button.click(
-        fn = lambda image: image,
-        inputs = [result],
-        outputs = [input_image]
+        fn=infer,
+        inputs=[prompt, input_images, mode_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
+        outputs=[result, seed]
     )
 
-demo.launch(mcp_server=True)
+demo.launch()

diffusers.zip

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

optimization.py

@@ -4,23 +4,35 @@
 from typing import Any
 from typing import Callable
 from typing import ParamSpec
-
 import spaces
 import torch
-from torch.utils._pytree import tree_map_only
-
-from optimization_utils import capture_component_call
-from optimization_utils import aoti_compile
-
+from spaces.zero.torch.aoti import ZeroGPUCompiledModel
+from spaces.zero.torch.aoti import ZeroGPUWeights
+from torch.utils._pytree import tree_map
 
 P = ParamSpec('P')
 
-
-TRANSFORMER_HIDDEN_DIM = torch.export.Dim('hidden', min=4096, max=8212)
+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': {1: TRANSFORMER_HIDDEN_DIM},
-    'img_ids': {0: TRANSFORMER_HIDDEN_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},
+        ),
+    },
 }
 
 INDUCTOR_CONFIGS = {
@@ -32,29 +44,34 @@ INDUCTOR_CONFIGS = {
     'triton.cudagraphs': True,
 }
 
-
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
 
-    @spaces.GPU(duration=1500)
-    def compile_transformer():
+    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
-
-        pipeline.transformer.fuse_qkv_projections()
+        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES[blocks_kind]
 
-        exported = torch.export.export(
-            mod=pipeline.transformer,
-            args=call.args,
-            kwargs=call.kwargs,
-            dynamic_shapes=dynamic_shapes,
-        )
+        with torch.no_grad():      
+            exported = torch.export.export(
+                mod=block,
+                args=call.args,
+                kwargs=call.kwargs,
+                dynamic_shapes=dynamic_shapes,
+            )
 
-        return aoti_compile(exported, INDUCTOR_CONFIGS)
+        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS).archive_file
 
-    transformer_config = pipeline.transformer.config
-    pipeline.transformer = compile_transformer()
-    pipeline.transformer.config = transformer_config # pyright: ignore[reportAttributeAccessIssue]
+    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,5 +1,8 @@
+git+https://github.com/huggingface/diffusers.git@flux2-klein
 transformers
-git+https://github.com/huggingface/diffusers.git
 accelerate
 safetensors
-sentencepiece
+bitsandbytes
+torchao
+kernels
+spaces==0.43.0