Upload 8 files

README.md

@@ -1,12 +1,14 @@
 ---
-title: Wan 2 2 First Last Frame
-emoji: 💻
-colorFrom: purple
+title: FLUX.1 Kontext
+emoji: ⚡
+colorFrom: green
 colorTo: gray
 sdk: gradio
 sdk_version: 5.29.1
 app_file: app.py
-pinned: false
+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,8 +1,8 @@
-import os
 # 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')
 
-# --- 1. Model Download and Setup (Diffusers Backend) ---
+# Actual demo code
 try:
     import spaces
 except:
@@ -12,662 +12,268 @@ except:
                 return lambda *dummy_args, **dummy_kwargs: function(*dummy_args, **dummy_kwargs)
             return decorator
 
-import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
-from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
-from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
-from diffusers.utils.export_utils import export_to_video
 import gradio as gr
-import imageio_ffmpeg
-import tempfile
-import shutil
-import subprocess
-import time
-from datetime import datetime
 import numpy as np
-from PIL import Image
+import torch
 import random
-import math
-import traceback
-import gc
-from gradio_client import Client, handle_file # Import for API call
-import zipfile
-
-# Import optimization and access compiled artifacts
-import optimization
+from datetime import datetime
 
-# Import the optimization function from the separate file
-from optimization import optimize_pipeline_
+from PIL import Image
+import tempfile
+import shutil
+from pathlib import Path
 
-# --- Constants and Model Loading ---
-MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+from diffusers import FluxKontextPipeline
+from diffusers.utils import load_image
 
-# --- NEW: Flexible Dimension Constants ---
-MAX_DIMENSION = 832
-MIN_DIMENSION = 480
-DIMENSION_MULTIPLE = 16
-SQUARE_SIZE = 480
+from optimization import optimize_pipeline_
 
 MAX_SEED = np.iinfo(np.int32).max
 
-FIXED_FPS = 24
-MIN_FRAMES_MODEL = 8
-MAX_FRAMES_MODEL = 81
-
-MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS, 1)
-MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS, 1)
+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]
-end_image_debug_value = [None]
 prompt_debug_value = [None]
-total_second_length_debug_value = [None]
-resolution_debug_value = [None]
-factor_debug_value = [None]
-allocation_time_debug_value = [None]
-
-default_negative_prompt = "Vibrant colors, overexposure, static, blurred details, subtitles, error, style, artwork, painting, image, still, overall gray, worst quality, low quality, JPEG compression residue, ugly, mutilated, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, malformed limbs, fused fingers, still image, cluttered background, three legs, many people in the background, walking backwards, overexposure, jumpcut, crossfader, "
-
-transformer = WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
-        subfolder='transformer',
-        torch_dtype=torch.bfloat16,
-        device_map='cuda',
-    )
-
-transformer_2 = WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
-        subfolder='transformer_2',
-        torch_dtype=torch.bfloat16,
-        device_map='cuda',
-    )
-
-pipe = WanImageToVideoPipeline.from_pretrained(
-    MODEL_ID,
-    transformer = transformer,
-    transformer_2 = transformer_2,
-    torch_dtype=torch.bfloat16,
-)
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config, shift=8.0)
-pipe.to('cuda')
-
-for i in range(3):
-    gc.collect()
-    torch.cuda.synchronize()
-    torch.cuda.empty_cache()
-
-optimize_pipeline_(pipe,
-    image=Image.new('RGB', (MAX_DIMENSION, MIN_DIMENSION)), 
-    prompt='prompt',
-    height=MIN_DIMENSION,
-    width=MAX_DIMENSION,
-    num_frames=MAX_FRAMES_MODEL,
-)
-
-def _escape_html(s: str) -> str:
-    return (s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;"))
-
-def _error_to_html(err: BaseException) -> str:
-    tb = traceback.format_exc()
-    return (
-        "<div style='padding:12px;border:1px solid #ff4d4f;background:#fff1f0;color:#a8071a;border-radius:8px;'>"
-        "<b>Generation failed</b><br/>"
-        f"<b>{_escape_html(type(err).__name__)}</b>: {_escape_html(str(err))}"
-        "<details style='margin-top:8px;'>"
-        "<summary>Show traceback</summary>"
-        f"<pre style='white-space:pre-wrap;margin-top:8px;'>{_escape_html(tb)}</pre>"
-        "</details>"
-        "</div>"
-    )
-
-# 20250508 pftq: for saving prompt to mp4 metadata comments
-def set_mp4_comments_imageio_ffmpeg(input_file, comments):
-    try:
-        # Get the path to the bundled FFmpeg binary from imageio-ffmpeg
-        ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe()
-
-        # Check if input file exists
-        if not os.path.exists(input_file):
-            #print(f"Error: Input file {input_file} does not exist")
-            return False
-
-        # Create a temporary file path
-        temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
-
-        # FFmpeg command using the bundled binary
-        command = [
-            ffmpeg_path,                   # Use imageio-ffmpeg's FFmpeg
-            '-i', input_file,              # input file
-            '-metadata', f'comment={comments}',  # set comment metadata
-            '-c:v', 'copy',                # copy video stream without re-encoding
-            '-c:a', 'copy',                # copy audio stream without re-encoding
-            '-y',                          # overwrite output file if it exists
-            temp_file                      # temporary output file
-        ]
-
-        # Run the FFmpeg command
-        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-
-        if result.returncode == 0:
-            # Replace the original file with the modified one
-            shutil.move(temp_file, input_file)
-            #print(f"Successfully added comments to {input_file}")
-            return True
-        else:
-            # Clean up temp file if FFmpeg fails
-            if os.path.exists(temp_file):
-                os.remove(temp_file)
-            #print(f"Error: FFmpeg failed with message:\n{result.stderr}")
-            return False
-
-    except Exception as e:
-        # Clean up temp file in case of other errors
-        if 'temp_file' in locals() and os.path.exists(temp_file):
-            os.remove(temp_file)
-        print(f"Error saving prompt to video metadata, ffmpeg may be required: "+str(e))
-        return False
-
-# --- 2. Image Processing and Application Logic ---
-def generate_end_frame(start_img, gen_prompt, progress=gr.Progress(track_tqdm=True)):
-    """Calls an external Gradio API to generate an image."""
-    if start_img is None:
-        raise gr.Error("Please provide a Start Frame first.")
-
-    hf_token = os.getenv("HF_TOKEN")
-    if not hf_token:
-        raise gr.Error("HF_TOKEN not found in environment variables. Please set it in your Space secrets.")
-
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
-        start_img.save(tmpfile.name)
-        tmp_path = tmpfile.name
-
-    progress(0.1, desc="Connecting to image generation API...")
-    client = Client("multimodalart/nano-banana-private")
-    
-    progress(0.5, desc=f"Generating with prompt: '{gen_prompt}'...")
-    try:
-        result = client.predict(
-            prompt=gen_prompt, 
-            images=[
-                {"image": handle_file(tmp_path)}
-            ],
-            manual_token=hf_token,
-            api_name="/unified_image_generator"
-        )
-    finally:
-        os.remove(tmp_path)
-
-    progress(1.0, desc="Done!")
-    print(result)
-    return result
-
-def switch_to_upload_tab():
-    """Returns a gr.Tabs update to switch to the first tab."""
-    return gr.Tabs(selected="upload_tab")
-
-
-def process_image_for_video(image: Image.Image, resolution: int) -> Image.Image:
+number_debug_value = [None]
+def save_on_path(img: Image, filename: str, format_: str = None) -> Path:
     """
-    Resizes an image based on the following rules for video generation.
+    Save `img` in a unique temporary folder under the given `filename`
+    and return its absolute path.
     """
-    width, height = image.size
-
-    if resolution < width * height:
-        scale = ((width * height) / resolution)**(.5)
-        new_width = width / scale
-        new_height = height / scale
-        final_width = int(math.floor(new_width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
-        final_height = int(math.floor(new_height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
-
-    elif width * height < (MIN_DIMENSION**2):
-        scale = ((MIN_DIMENSION**2) / (width * height))**(.5)
-        new_width = width * scale
-        new_height = height * scale
-        final_width = int(math.ceil(new_width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
-        final_height = int(math.ceil(new_height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
-
-    else:
-        final_width = int(round(width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
-        final_height = int(round(height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
-
-    return image.resize((final_width, final_height), Image.Resampling.LANCZOS)
-
-def resize_and_crop_to_match(target_image, reference_image):
-    """Resizes the target image to match the reference image's dimensions."""
-    ref_width, ref_height = reference_image.size
-    return target_image.resize((ref_width, ref_height), Image.Resampling.LANCZOS)
-
-def crop_to_match(target_image, reference_image):
-    """Resizes and center-crops the target image to match the reference image's dimensions."""
-    ref_width, ref_height = reference_image.size
-    target_width, target_height = target_image.size
-    scale = max(ref_width / target_width, ref_height / target_height)
-    new_width, new_height = int(target_width * scale), int(target_height * scale)
-    resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
-    left, top = (new_width - ref_width) // 2, (new_height - ref_height) // 2
-    return resized.crop((left, top, left + ref_width, top + ref_height))
-
-def init_view():
-    return gr.update(interactive = True)
-
-def output_video_change(output_video):
-    print('Log output: ' + str(output_video))
-    return [gr.update(visible = True)] * 2
+    # 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 generate_video(
-    start_image_pil,
-    end_image_pil,
+@spaces.GPU(duration=40)
+def infer(
+    input_image,
     prompt,
-    negative_prompt=default_negative_prompt,
-    resolution=500000,
-    duration_seconds=2.1,
-    steps=8,
-    guidance_scale=1,
-    guidance_scale_2=1,
-    seed=42,
-    randomize_seed=True,
+    seed = 42,
+    randomize_seed = False,
+    guidance_scale = 2.5,
+    steps = 28,
+    width = -1,
+    height = -1,
     progress=gr.Progress(track_tqdm=True)
 ):
-    start = time.time()
-    allocation_time = 120
-    factor = 1
+    """
+    Perform image editing using the FLUX.1 Kontext pipeline.
     
-    if input_image_debug_value[0] is not None or end_image_debug_value[0] is not None or prompt_debug_value[0] is not None or total_second_length_debug_value[0] is not None or allocation_time_debug_value[0] is not None or resolution_debug_value[0] is not None or factor_debug_value[0] is not None:
-        start_image_pil = input_image_debug_value[0]
-        end_image_pil = end_image_debug_value[0]
-        prompt = prompt_debug_value[0]
-        duration_seconds = total_second_length_debug_value[0]
-        resolution = resolution_debug_value[0]
-        factor = factor_debug_value[0]
-        allocation_time = allocation_time_debug_value[0]
-
-    if start_image_pil is None or end_image_pil is None:
-        raise gr.Error("Please upload both a start and an end image.")
-
-    # Step 1: Process the start image to get our target dimensions based on the new rules.
-    processed_start_image = process_image_for_video(start_image_pil, resolution)
+    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.
     
-    # Step 2: Make the end image match the *exact* dimensions of the processed start image.
-    processed_end_image = resize_and_crop_to_match(end_image_pil, processed_start_image)
+    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).
     
-    target_height, target_width = processed_start_image.height, processed_start_image.width
-
-    # Handle seed and frame count
-    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
-    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
-
-    progress(0.2, desc=f"Generating {num_frames} frames at {target_width}x{target_height} (seed: {current_seed})...")
-
-    progress(0.1, desc="Preprocessing images...")
-    print("Generate a video with the prompt: " + prompt)
-    output_frames_list = None
-    caught_error = None
-    while factor >= 1 and int(allocation_time) > 0:
-        try:
-            output_frames_list = generate_video_on_gpu(
-                start_image_pil,
-                end_image_pil,
-                prompt,
-                negative_prompt,
-                int(steps),
-                float(guidance_scale),
-                float(guidance_scale_2),
-                progress,
-                allocation_time,
-                target_height,
-                target_width,
-                current_seed,
-                (int(((num_frames * factor) - 1) / 4) * 4) + 1,
-                processed_start_image,
-                processed_end_image
-            )
-            factor = 0
-            caught_error = None
-        except BaseException as err:
-            print("An exception occurred: " + str(err))
-            caught_error = err
-            try:
-                print('e.message: ' + err.message) # No GPU is currently available for you after 60s
-            except Exception as e2:
-                print('Failure')
-            if not str(err).startswith("No GPU is currently available for you after 60s"):
-                factor -= .003
-                allocation_time = int(allocation_time) - 1
-        except:
-            print("An error occurred")
-            caught_error = None
-            if not str(e).startswith("No GPU is currently available for you after 60s"):
-                factor -= .003
-                allocation_time = int(allocation_time) - 1
-
-    if caught_error is not None:
-        return [gr.skip(), gr.skip(), gr.skip(), gr.update(value=_error_to_html(caught_error), visible=True), gr.skip()]
-
-    input_image_debug_value[0] = end_image_debug_value[0] = prompt_debug_value[0] = total_second_length_debug_value[0] = allocation_time_debug_value[0] = factor_debug_value[0] = None
+    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
     
-    progress(0.9, desc="Encoding and saving video...")
+    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)
     
-    video_path = 'wan_' + datetime.now().strftime("%Y-%m-%d_%H-%M-%S.%f") + '.mp4'
-
-    export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
-    set_mp4_comments_imageio_ffmpeg(video_path, f"Prompt: {prompt} | Negative Prompt: {negative_prompt}");
-    print("Video exported: " + video_path)
+    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]
     
-    progress(1.0, desc="Done!")
-    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 video 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 video has " + str(int(num_frames * factor)) + " frames. " + \
-    "The video resolution is " + str(target_width) + \
-    " pixels large and " + str(target_height) + \
-    " pixels high, so a resolution of " + f'{target_width * target_height:,}' + " pixels." + \
-    " Your prompt is saved into the metadata of the video."
-    return [video_path, gr.update(value = video_path, visible = True, interactive = True), current_seed, gr.update(value = information, visible = True), gr.update(interactive = False)]
-
-def get_duration(
-    start_image_pil,
-    end_image_pil,
-    prompt,
-    negative_prompt,
-    steps,
-    guidance_scale,
-    guidance_scale_2,
-    progress,
-    allocation_time,
-    target_height,
-    target_width,
-    current_seed,
-    num_frames,
-    processed_start_image,
-    processed_end_image
-):
-    return allocation_time
-
-@torch.no_grad()
-@spaces.GPU(duration=get_duration)
-def generate_video_on_gpu(
-    start_image_pil,
-    end_image_pil,
-    prompt,
-    negative_prompt,
-    steps,
-    guidance_scale,
-    guidance_scale_2,
-    progress,
-    allocation_time,
-    target_height,
-    target_width,
-    current_seed,
-    num_frames,
-    processed_start_image,
-    processed_end_image
-):
-    """
-    Generates a video by interpolating between a start and end image, guided by a text prompt,
-    using the diffusers Wan2.2 pipeline.
-    """
-
-    output_frames_list = pipe(
-        image=processed_start_image,
-        last_image=processed_end_image,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        height=target_height,
-        width=target_width,
-        num_frames=num_frames,
-        guidance_scale=guidance_scale,
-        guidance_scale_2=guidance_scale_2,
-        num_inference_steps=steps,
-        generator=torch.Generator(device="cuda").manual_seed(current_seed),
-    ).frames[0]
-
-    return output_frames_list
-
-def export_compiled_transformers_to_zip() -> str:
-    """
-    Bundle compiled_transformer_1 and compiled_transformer_2 into a zip file and return the file path.
-    """
-    ct1 = getattr(optimization, "COMPILED_TRANSFORMER_1", None)
-    ct2 = getattr(optimization, "COMPILED_TRANSFORMER_2", None)
-
-    if ct1 is None or ct2 is None:
-        raise gr.Error("Compiled transformers are not available yet (compilation may have failed).")
-
-    payload_1 = ct1.to_serializable_dict()
-    payload_2 = ct2.to_serializable_dict()
-
-    tmp_zip = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
-    tmp_zip.close()
-
-    with zipfile.ZipFile(tmp_zip.name, "w", compression=zipfile.ZIP_DEFLATED) as zf:
-        # store with torch.save so users can load easily with torch.load()
-        buf1 = tempfile.NamedTemporaryFile(suffix=".pt", delete=False)
-        buf1.close()
-        torch.save(payload_1, buf1.name)
-
-        buf2 = tempfile.NamedTemporaryFile(suffix=".pt", delete=False)
-        buf2.close()
-        torch.save(payload_2, buf2.name)
-
-        zf.write(buf1.name, arcname="compiled_transformer_1.pt")
-        zf.write(buf2.name, arcname="compiled_transformer_2.pt")
-
-    # cleanup intermediate .pt
+    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:
-        os.remove(buf1.name)
-        os.remove(buf2.name)
+        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:
-        pass
-
-    return tmp_zip.name
-
-
-# --- 3. Gradio User Interface ---
-
+        print("Error")
+    return gallery, seed
 
-
-js = """
-function createGradioAnimation() {
-    window.addEventListener("beforeunload", function(e) {
-        if (document.getElementById('dummy_button_id') && !document.getElementById('dummy_button_id').disabled) {
-            var confirmationMessage = 'A process is still running. '
-                                    + 'If you leave before saving, your changes will be lost.';
-    
-            (e || window.event).returnValue = confirmationMessage;
-        }
-        return confirmationMessage;
-    });
-    return 'Animation created';
+css="""
+#col-container {
+    margin: 0 auto;
+    max-width: 960px;
 }
 """
 
-# Gradio interface
-with gr.Blocks(js=js) as app:
-    gr.Markdown("# Wan 2.2 First/Last Frame Video Fast")
-    gr.Markdown("Based on the [Wan 2.2 First/Last Frame workflow](https://www.reddit.com/r/StableDiffusion/comments/1me4306/psa_wan_22_does_first_frame_last_frame_out_of_the/), applied to 🧨 Diffusers + [lightx2v/Wan2.2-Lightning](https://huggingface.co/lightx2v/Wan2.2-Lightning) 8-step LoRA")
-
-    with gr.Row(elem_id="general_items"):
-        with gr.Column():
-            with gr.Group(elem_id="group_all"):
+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():
-                    start_image = gr.Image(type="pil", label="Start Frame", sources=["upload", "clipboard"])
-                    # Capture the Tabs component in a variable and assign IDs to tabs
-                    with gr.Tabs(elem_id="group_tabs") as tabs:
-                        with gr.TabItem("Upload", id="upload_tab"):
-                            end_image = gr.Image(type="pil", label="End Frame", sources=["upload", "clipboard"])
-                        with gr.TabItem("Generate", id="generate_tab"):
-                            generate_5seconds = gr.Button("Generate scene 5 seconds in the future", elem_id="fivesec")
-                            gr.Markdown("Generate a custom end-frame with an edit model like [Nano Banana](https://huggingface.co/spaces/multimodalart/nano-banana) or [Qwen Image Edit](https://huggingface.co/spaces/multimodalart/Qwen-Image-Edit-Fast)", elem_id="or_item")
-                prompt = gr.Textbox(label="Prompt", info="Describe the transition between the two images", placeholder="The creature starts to move")
-
+                    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):
-                    duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=2.1, label="Video Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
-                    negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
-                    resolution = gr.Dropdown([
-                        ["400,000 px (working)", 400000],
-                        ["465,920 px (working)", 465920],
-                        ["495,616 px (working)", 495616],
-                        ["500,000 px (working)", 500000],
-                        ["600,000 px (working)", 600000],
-                        ["700,000 px (working)", 700000],
-                        ["800,000 px (working)", 800000],
-                        ["900,000 px (working)", 900000],
-                        ["1,000,000 px (working)", 1000000],
-                        ["1,100,000 px (untested)", 1100000],
-                        ["1,200,000 px (untested)", 1200000],
-                        ["1,300,000 px (untested)", 1300000],
-                        ["1,400,000 px (untested)", 1400000],
-                        ["1,500,000 px (untested)", 1500000]
-                    ], value=465920, label="Resolution (width x height)", info="Less if the image is smaller")
-                    steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=8, label="Inference Steps")
-                    guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - high noise")
-                    guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - low noise")
-                    with gr.Row():
-                        seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
-                        randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True)
-
-                generate_button = gr.Button("🚀 Generate Video", variant="primary")
-                dummy_button = gr.Button(elem_id = "dummy_button_id", visible = False, interactive = False)
-
-        with gr.Column():
-            output_video = gr.Video(label="Generated Video", autoplay = True, loop = True)
-            download_button = gr.DownloadButton(elem_id="download_btn", interactive = True)
-            video_information = gr.HTML(value = "")
-
-    with gr.Accordion("🔧 Compilation artifacts (advanced)", open=False):
-        gr.Markdown(
-            "Télécharge les artefacts compilés AOTInductor générés au démarrage (transformer + transformer_2)."
-        )
-        export_btn = gr.Button("📦 Préparer l'archive des transformers compilés")
-        compiled_download = gr.DownloadButton(label="⬇️ Télécharger compiled_transformers.zip", interactive=False)
-
-        def _build_and_enable_download():
-            path = export_compiled_transformers_to_zip()
-            return gr.update(value=path, interactive=True)
-
-        export_btn.click(fn=_build_and_enable_download, inputs=None, outputs=compiled_download)
-
-    # Main video generation button
-    ui_inputs = [
-        start_image,
-        end_image,
-        prompt,
-        negative_prompt_input,
-        resolution,
-        duration_seconds_input,
-        steps_slider,
-        guidance_scale_input,
-        guidance_scale_2_input,
-        seed_input,
-        randomize_seed_checkbox
-    ]
-    ui_outputs = [output_video, download_button, seed_input, video_information, dummy_button]
-
-    generate_button.click(fn = init_view, inputs = [], outputs = [dummy_button], queue = False, show_progress = False).success(
-        fn = generate_video,
-        inputs = ui_inputs,
-        outputs = ui_outputs
-    )
-
-    generate_5seconds.click(
-        fn=switch_to_upload_tab,
-        inputs=None,
-        outputs=[tabs]
-    ).then( 
-        fn=lambda img: generate_end_frame(img, "this image is a still frame from a movie. generate a new frame with what happens on this scene 5 seconds in the future"),
-        inputs=[start_image],
-        outputs=[end_image]
-    ).success(
-        fn=generate_video,
-        inputs=ui_inputs,
-        outputs=ui_outputs
-    )
-    
-    output_video.change(
-        fn=output_video_change,
-        inputs=[output_video],
-        outputs=[download_button, video_information],
-        js="document.getElementById('download_btn').click()"
-    )
+                    
+                    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):
+        result_gallery = gr.Gallery(label = 'Downloadable results', show_label = True, interactive = False, elem_id = "gallery1")
         gr.Examples(
-            examples=[["Schoolboy_without_backpack.webp", "Schoolboy_with_backpack.webp", "The schoolboy puts on his schoolbag."]],
-            inputs=[start_image, end_image, prompt],
-            outputs=ui_outputs,
-            fn=generate_video,
+            examples=[
+                ["monster.png", "Make this monster ride a skateboard on the beach"]
+            ],
+            inputs=[input_image, prompt],
+            outputs=[result_gallery, seed],
+            fn=infer_example,
             run_on_click=True,
             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")
-        end_image_debug=gr.Image(type="pil", label="End Image Debug")
-        total_second_length_debug=gr.Slider(label="Duration Debug", minimum=1, maximum=120, value=5, step=0.1)
-        resolution_debug = gr.Dropdown([
-            ["400,000 px", 400000],
-            ["465,920 px", 465920],
-            ["495,616 px", 495616],
-            ["500,000 px", 500000],
-            ["600,000 px", 600000],
-            ["700,000 px", 700000],
-            ["800,000 px", 800000],
-            ["900,000 px", 900000],
-            ["1,000,000 px", 1000000],
-            ["1,100,000 px", 1100000],
-            ["1,200,000 px", 1200000],
-            ["1,300,000 px", 1300000],
-            ["1,400,000 px", 1400000],
-            ["1,500,000 px", 1500000]
-        ], value=500000, label="Resolution Debug")
-        factor_debug=gr.Slider(label="Factor Debug", minimum=1, maximum=100, value=3.2, step=0.1)
-        allocation_time_debug=gr.Slider(label="Allocation Debug", minimum=1, maximum=60 * 20, value=720, step=1)
+        number_debug=gr.Slider(label="Number Debug", minimum=1, maximum=50, step=1, value=50)
 
-    def handle_field_debug_change(
-        input_image_debug_data,
-        end_image_debug_data,
-        prompt_debug_data,
-        total_second_length_debug_data,
-        resolution_debug_data,
-        factor_debug_data,
-        allocation_time_debug_data
-    ):
-        input_image_debug_value[0] = input_image_debug_data
-        end_image_debug_value[0] = end_image_debug_data
+    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
-        total_second_length_debug_value[0] = total_second_length_debug_data
-        resolution_debug_value[0] = resolution_debug_data
-        factor_debug_value[0] = factor_debug_data
-        allocation_time_debug_value[0] = allocation_time_debug_data
+        input_image_debug_value[0] = input_image_debug_data
+        number_debug_value[0] = number_debug_data
         return []
 
-    inputs_debug=[input_image_debug, end_image_debug, prompt_debug, total_second_length_debug, resolution_debug, factor_debug, allocation_time_debug]
+    inputs_debug=[input_image_debug, prompt_debug, number_debug]
 
     input_image_debug.upload(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-    end_image_debug.upload(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
     prompt_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-    total_second_length_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-    resolution_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-    factor_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-    allocation_time_debug.change(fn=handle_field_debug_change, inputs=inputs_debug, outputs=[])
-    
-    gr.Examples(
-        label = "Examples from demo",
-        examples = [
-            ["poli_tower.png", "tower_takes_off.png", "The man turns around."],
-            ["ugly_sonic.jpeg", "squatting_sonic.png", "पात्रं क्षेपणास्त्रं चकमाति।"],
-            ["Schoolboy_without_backpack.webp", "Schoolboy_with_backpack.webp", "The schoolboy puts on his schoolbag."],
-        ],
-        inputs = [start_image, end_image, prompt],
-        outputs = ui_outputs,
-        fn = generate_video,
-        cache_examples = False,
+    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]
     )
 
-if __name__ == "__main__":
-    app.launch(mcp_server=True, share=True)
+demo.launch(mcp_server=True)

optimization.py

@@ -1,156 +1,60 @@
-"""
-"""
-
-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
-
-
-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_DYNAMIC_SHAPES = {
-    'hidden_states': {
-        2: LATENT_FRAMES_DIM,
-        3: 2 * LATENT_PATCHED_HEIGHT_DIM,
-        4: 2 * LATENT_PATCHED_WIDTH_DIM,
-    },
-}
-
-INDUCTOR_CONFIGS = {
-    'conv_1x1_as_mm': True,
-    'epilogue_fusion': False,
-    'coordinate_descent_tuning': True,
-    'coordinate_descent_check_all_directions': True,
-    'max_autotune': True,
-    '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()
-
-        with capture_component_call(pipeline, "transformer") 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,
-        )
-
-        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",
-            )
-        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)
-
-    pipeline.transformer_2.forward = compiled_transformer_2
-    drain_module_parameters(pipeline.transformer_2)
+"""
+"""
+
+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
+
+
+P = ParamSpec('P')
+
+
+TRANSFORMER_HIDDEN_DIM = torch.export.Dim('hidden', min=4096, max=8212)
+
+TRANSFORMER_DYNAMIC_SHAPES = {
+    'hidden_states': {1: TRANSFORMER_HIDDEN_DIM},
+    'img_ids': {0: TRANSFORMER_HIDDEN_DIM},
+}
+
+INDUCTOR_CONFIGS = {
+    'conv_1x1_as_mm': True,
+    'epilogue_fusion': False,
+    'coordinate_descent_tuning': True,
+    'coordinate_descent_check_all_directions': True,
+    'max_autotune': True,
+    'triton.cudagraphs': True,
+}
+
+
+def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
+
+    @spaces.GPU(duration=1500)
+    def compile_transformer():
+
+        with capture_component_call(pipeline, 'transformer') 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()
+
+        exported = torch.export.export(
+            mod=pipeline.transformer,
+            args=call.args,
+            kwargs=call.kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+
+        return aoti_compile(exported, INDUCTOR_CONFIGS)
+
+    transformer_config = pipeline.transformer.config
+    pipeline.transformer = compile_transformer()
+    pipeline.transformer.config = transformer_config # pyright: ignore[reportAttributeAccessIssue]

optimization_utils.py

@@ -1,166 +1,96 @@
-"""
-"""
-import contextlib
-from contextvars import ContextVar
-from io import BytesIO
-from typing import Any
-from typing import cast
-from unittest.mock import patch
-
-import torch
-from torch._inductor.package.package import package_aoti
-from torch.export.pt2_archive._package import AOTICompiledModel
-from torch.export.pt2_archive._package_weights import Weights
-
-
-INDUCTOR_CONFIGS_OVERRIDES = {
-    'aot_inductor.package_constants_in_so': False,
-    'aot_inductor.package_constants_on_disk': True,
-    'aot_inductor.package': True,
-}
-
-
-class ZeroGPUWeights:
-    def __init__(self, constants_map: dict[str, torch.Tensor], to_cuda: bool = False):
-        if to_cuda:
-            self.constants_map = {name: tensor.to('cuda') for name, tensor in constants_map.items()}
-        else:
-            self.constants_map = constants_map
-
-    def __reduce__(self):
-        constants_map: dict[str, torch.Tensor] = {}
-        for name, tensor in self.constants_map.items():
-            tensor_ = torch.empty_like(tensor, device='cpu').pin_memory()
-            constants_map[name] = tensor_.copy_(tensor).detach().share_memory_()
-        return ZeroGPUWeights, (constants_map, True)
-
-
-class ZeroGPUCompiledModel:
-    def __init__(self, archive_file: torch.types.FileLike, weights: ZeroGPUWeights):
-        self.archive_file = archive_file
-        self.weights = weights
-        self.compiled_model: ContextVar[AOTICompiledModel | None] = ContextVar('compiled_model', default=None)
-
-    def __call__(self, *args, **kwargs):
-        if (compiled_model := self.compiled_model.get()) is None:
-            compiled_model = cast(AOTICompiledModel, torch._inductor.aoti_load_package(self.archive_file))
-            compiled_model.load_constants(self.weights.constants_map, check_full_update=True, user_managed=True)
-            self.compiled_model.set(compiled_model)
-        return compiled_model(*args, **kwargs)
-
-    def __reduce__(self):
-        return ZeroGPUCompiledModel, (self.archive_file, self.weights)
-
-    # --- NEW: stable payload export (matches your zip export format) ---
-    def to_serializable_dict(self) -> dict[str, Any]:
-        """
-        Stable representation that can be stored with torch.save and re-loaded later,
-        without depending on runtime state.
-        """
-        if hasattr(self.archive_file, "getvalue"):
-            archive_bytes = self.archive_file.getvalue()
-        else:
-            pos = self.archive_file.tell()
-            self.archive_file.seek(0)
-            archive_bytes = self.archive_file.read()
-            self.archive_file.seek(pos)
-
-        constants_cpu = {k: v.detach().to("cpu") for k, v in self.weights.constants_map.items()}
-        return {
-            "format": "zerogpu_aoti_v1",
-            "archive_bytes": archive_bytes,
-            "constants_map": constants_cpu,
-        }
-
-
-# --- NEW: loader for the payload dict saved in .pt ---
-class ZeroGPUCompiledModelFromDict:
-    """
-    Rebuilds a callable AOTI model from a dict payload produced by ZeroGPUCompiledModel.to_serializable_dict().
-
-    Important: constants are moved to CUDA on first call to avoid:
-      - dtype/device mismatches (CPU bf16 weights vs CUDA bf16 inputs),
-      - black outputs due to wrong device handling.
-    """
-    def __init__(self, payload: dict[str, Any], device: str = "cuda"):
-        if payload.get("format") != "zerogpu_aoti_v1":
-            raise ValueError(f"Unsupported payload format: {payload.get('format')}")
-        self.archive_file = BytesIO(payload["archive_bytes"])
-        constants = payload["constants_map"]
-        constants = {k: v.to(device=device, dtype=torch.bfloat16).contiguous() for k, v in constants.items()}
-        payload["constants_map"] = constants
-        self.constants_map_cpu: dict[str, torch.Tensor] = payload["constants_map"]
-        self.device = device
-        self.compiled_model: ContextVar[AOTICompiledModel | None] = ContextVar("compiled_model", default=None)
-        self._loaded_constants = False
-
-    def __call__(self, *args, **kwargs):
-        if (compiled_model := self.compiled_model.get()) is None:
-            compiled_model = cast(AOTICompiledModel, torch._inductor.aoti_load_package(self.archive_file))
-            self.compiled_model.set(compiled_model)
-
-        if not self._loaded_constants:
-            # Move constants to target device (cuda) and keep dtype as-is (bf16)
-            constants_map = {k: v.to(device="cuda", dtype=torch.bfloat16).contiguous() for k, v in self.constants_map_cpu.items()}
-            compiled_model.load_constants(constants_map, check_full_update=True, user_managed=True)
-            self._loaded_constants = True
-
-        return compiled_model(*args, **kwargs)
-
-
-def aoti_compile(
-    exported_program: torch.export.ExportedProgram,
-    inductor_configs: dict[str, Any] | None = None,
-):
-    inductor_configs = (inductor_configs or {}) | INDUCTOR_CONFIGS_OVERRIDES
-    gm = cast(torch.fx.GraphModule, exported_program.module())
-    assert exported_program.example_inputs is not None
-    args, kwargs = exported_program.example_inputs
-    artifacts = torch._inductor.aot_compile(gm, args, kwargs, options=inductor_configs)
-    archive_file = BytesIO()
-    files: list[str | Weights] = [file for file in artifacts if isinstance(file, str)]
-    package_aoti(archive_file, files)
-    weights, = (artifact for artifact in artifacts if isinstance(artifact, Weights))
-    zerogpu_weights = ZeroGPUWeights({name: weights.get_weight(name)[0] for name in weights})
-    return ZeroGPUCompiledModel(archive_file, zerogpu_weights)
-
-
-@contextlib.contextmanager
-def capture_component_call(
-    pipeline: Any,
-    component_name: str,
-    component_method='forward',
-):
-
-    class CapturedCallException(Exception):
-        def __init__(self, *args, **kwargs):
-            super().__init__()
-            self.args = args
-            self.kwargs = kwargs
-
-    class CapturedCall:
-        def __init__(self):
-            self.args: tuple[Any, ...] = ()
-            self.kwargs: dict[str, Any] = {}
-
-    component = getattr(pipeline, component_name)
-    captured_call = CapturedCall()
-
-    def capture_call(*args, **kwargs):
-        raise CapturedCallException(*args, **kwargs)
-
-    with patch.object(component, component_method, new=capture_call):
-        try:
-            yield captured_call
-        except CapturedCallException as e:
-            captured_call.args = e.args
-            captured_call.kwargs = e.kwargs
-
-
-def drain_module_parameters(module: torch.nn.Module):
-    state_dict_meta = {name: {'device': tensor.device, 'dtype': tensor.dtype} for name, tensor in module.state_dict().items()}
-    state_dict = {name: torch.nn.Parameter(torch.empty_like(tensor, device='cpu')) for name, tensor in module.state_dict().items()}
-    module.load_state_dict(state_dict, assign=True)
-    for name, param in state_dict.items():
-        meta = state_dict_meta[name]
-        param.data = torch.Tensor([]).to(**meta)
+"""
+"""
+import contextlib
+from contextvars import ContextVar
+from io import BytesIO
+from typing import Any
+from typing import cast
+from unittest.mock import patch
+
+import torch
+from torch._inductor.package.package import package_aoti
+from torch.export.pt2_archive._package import AOTICompiledModel
+from torch.export.pt2_archive._package_weights import TensorProperties
+from torch.export.pt2_archive._package_weights import Weights
+
+
+INDUCTOR_CONFIGS_OVERRIDES = {
+    'aot_inductor.package_constants_in_so': False,
+    'aot_inductor.package_constants_on_disk': True,
+    'aot_inductor.package': True,
+}
+
+
+class ZeroGPUCompiledModel:
+    def __init__(self, archive_file: torch.types.FileLike, weights: Weights, cuda: bool = False):
+        self.archive_file = archive_file
+        self.weights = weights
+        if cuda:
+            self.weights_to_cuda_()
+        self.compiled_model: ContextVar[AOTICompiledModel | None] = ContextVar('compiled_model', default=None)
+    def weights_to_cuda_(self):
+        for name in self.weights:
+            tensor, properties = self.weights.get_weight(name)
+            self.weights[name] = (tensor.to('cuda'), properties)
+    def __call__(self, *args, **kwargs):
+        if (compiled_model := self.compiled_model.get()) is None:
+            constants_map = {name: value[0] for name, value in self.weights.items()}
+            compiled_model = cast(AOTICompiledModel, torch._inductor.aoti_load_package(self.archive_file))
+            compiled_model.load_constants(constants_map, check_full_update=True, user_managed=True)
+            self.compiled_model.set(compiled_model)
+        return compiled_model(*args, **kwargs)
+    def __reduce__(self):
+        weight_dict: dict[str, tuple[torch.Tensor, TensorProperties]] = {}
+        for name in self.weights:
+            tensor, properties = self.weights.get_weight(name)
+            tensor_ = torch.empty_like(tensor, device='cpu').pin_memory()
+            weight_dict[name] = (tensor_.copy_(tensor).detach().share_memory_(), properties)
+        return ZeroGPUCompiledModel, (self.archive_file, Weights(weight_dict), True)
+
+
+def aoti_compile(
+    exported_program: torch.export.ExportedProgram,
+    inductor_configs: dict[str, Any] | None = None,
+):
+    inductor_configs = (inductor_configs or {}) | INDUCTOR_CONFIGS_OVERRIDES
+    gm = cast(torch.fx.GraphModule, exported_program.module())
+    assert exported_program.example_inputs is not None
+    args, kwargs = exported_program.example_inputs
+    artifacts = torch._inductor.aot_compile(gm, args, kwargs, options=inductor_configs)
+    archive_file = BytesIO()
+    files: list[str | Weights] = [file for file in artifacts if isinstance(file, str)]
+    package_aoti(archive_file, files)
+    weights, = (artifact for artifact in artifacts if isinstance(artifact, Weights))
+    return ZeroGPUCompiledModel(archive_file, weights)
+
+
+@contextlib.contextmanager
+def capture_component_call(
+    pipeline: Any,
+    component_name: str,
+    component_method='forward',
+):
+
+    class CapturedCallException(Exception):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+            self.args = args
+            self.kwargs = kwargs
+
+    class CapturedCall:
+        def __init__(self):
+            self.args: tuple[Any, ...] = ()
+            self.kwargs: dict[str, Any] = {}
+
+    component = getattr(pipeline, component_name)
+    captured_call = CapturedCall()
+
+    def capture_call(*args, **kwargs):
+        raise CapturedCallException(*args, **kwargs)
+
+    with patch.object(component, component_method, new=capture_call):
+        try:
+            yield captured_call
+        except CapturedCallException as e:
+            captured_call.args = e.args
+            captured_call.kwargs = e.kwargs

requirements.txt

@@ -1,12 +1,5 @@
-git+https://github.com/YassineT-cdc/diffusers.git@wan22-loras_optimized_contigous
-
-transformers==4.57.3
-accelerate==1.12.0
-safetensors==0.7.0
-sentencepiece==0.2.1
-peft==0.18.0
-ftfy==6.3.1
-imageio==2.37.2
-imageio-ffmpeg==0.6.0
-
-torchao==0.14.1
+transformers
+git+https://github.com/huggingface/diffusers.git
+accelerate
+safetensors
+sentencepiece