Upload 8 files

README.md

@@ -1,12 +1,14 @@
----
-title: FLUX.2 [Klein] 4B
-emoji: 💻
-colorFrom: blue
-colorTo: gray
-sdk: gradio
-sdk_version: 5.29.1
-app_file: app.py
-pinned: true
----
-
+---
+title: FLUX.1 Kontext
+emoji: ⚡
+colorFrom: green
+colorTo: gray
+sdk: gradio
+sdk_version: 5.29.1
+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,464 +1,299 @@
-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 = []
-    try:
-        for i in range(number):
-            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:
-        print("Error")
-    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()
+# 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 gradio as gr
+import numpy as np
+import torch
+import random
+import os
+from datetime import datetime
+
+from PIL import Image
+import tempfile
+import zipfile
+import shutil
+from pathlib import Path
+
+from diffusers import FluxKontextPipeline
+from diffusers.utils import load_image
+
+from optimization import optimize_pipeline_
+
+MAX_SEED = np.iinfo(np.int32).max
+
+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_"))
+    
+    # 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
+
+@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.
+    
+    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.
+    
+    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).
+    
+    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
+    
+    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]
+    
+    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="""
+#col-container {
+    margin: 0 auto;
+    max-width: 960px;
+}
+"""
+
+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)
+        """)
+        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')",
+                        container=False,
+                    )
+                    run_button = gr.Button(value="🚀 Edit", variant = "primary", scale=0)
+                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)
+                    
+                    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.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(interactive = False, elem_id = "gallery1")
+        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,
+            cache_examples=True,
+            cache_mode='lazy'
+        )
+        prompt_debug=gr.Textbox()
+        input_image_debug=gr.Image(type="pil")
+        number_debug=gr.Slider(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 []
+
+    inputs_debug=[input_image_debug, prompt_debug, number_debug]
+
+    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]
+    )
+
+demo.launch(mcp_server=True)

optimization.py

@@ -4,35 +4,23 @@
 from typing import Any
 from typing import Callable
 from typing import ParamSpec
+
 import spaces
 import torch
-from spaces.zero.torch.aoti import ZeroGPUCompiledModel
-from spaces.zero.torch.aoti import ZeroGPUWeights
-from torch.utils._pytree import tree_map
+from torch.utils._pytree import tree_map_only
+
+from optimization_utils import capture_component_call
+from optimization_utils import aoti_compile
+
 
 P = ParamSpec('P')
 
-TRANSFORMER_IMAGE_DIM = torch.export.Dim('image_seq_length', min=4096, max=16384) # min: 0 images, max: 3 (1024x1024) images
+
+TRANSFORMER_HIDDEN_DIM = torch.export.Dim('hidden', min=4096, max=8212)
 
 TRANSFORMER_DYNAMIC_SHAPES = {
-    '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},
-        ),
-    },
+    'hidden_states': {1: TRANSFORMER_HIDDEN_DIM},
+    'img_ids': {0: TRANSFORMER_HIDDEN_DIM},
 }
 
 INDUCTOR_CONFIGS = {
@@ -44,34 +32,29 @@ INDUCTOR_CONFIGS = {
     'triton.cudagraphs': True,
 }
 
+
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
 
-    blocks = {
-        'double': pipeline.transformer.transformer_blocks,
-        'single': pipeline.transformer.single_transformer_blocks,
-    }
+    @spaces.GPU(duration=1500)
+    def compile_transformer():
 
-    @spaces.GPU(duration=1200)
-    def compile_block(blocks_kind: str):
-        block = blocks[blocks_kind][0]
-        with spaces.aoti_capture(block) as call:
+        with capture_component_call(pipeline, 'transformer') as call:
             pipeline(*args, **kwargs)
 
-        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
-        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES[blocks_kind]
+        dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda t: None, call.kwargs)
+        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
+
+        pipeline.transformer.fuse_qkv_projections()
 
-        with torch.no_grad():      
-            exported = torch.export.export(
-                mod=block,
-                args=call.args,
-                kwargs=call.kwargs,
-                dynamic_shapes=dynamic_shapes,
-            )
+        exported = torch.export.export(
+            mod=pipeline.transformer,
+            args=call.args,
+            kwargs=call.kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
 
-        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS).archive_file
+        return aoti_compile(exported, INDUCTOR_CONFIGS)
 
-    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)
+    transformer_config = pipeline.transformer.config
+    pipeline.transformer = compile_transformer()
+    pipeline.transformer.config = transformer_config # pyright: ignore[reportAttributeAccessIssue]

requirements.txt

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