Initial commit with LFS

.gitattributes

@@ -0,0 +1,43 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+cat.png filter=lfs diff=lfs merge=lfs -text
+flowers.png filter=lfs diff=lfs merge=lfs -text
+monster.png filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: FLUX.1 Kontext
+emoji: ⚡
+colorFrom: green
+colorTo: gray
+sdk: gradio
+sdk_version: 5.34.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

@@ -0,0 +1,345 @@
+# 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
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import random
+from PIL import Image, ImageOps
+
+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("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
+pipe.load_lora_weights("ovi054/Draw2Photo")
+pipe.fuse_lora()
+# optimize_pipeline_(pipe, image=Image.new("RGB", (512, 512)), prompt='prompt')
+
+
+def add_overlay(base_img, overlay_img, margin=20):
+    """
+    Pastes an overlay image onto the top-right corner of a base image.
+
+    The overlay is resized to be 1/5th of the width of the base image,
+    maintaining its aspect ratio.
+
+    Args:
+        base_img (PIL.Image.Image): The main image.
+        overlay_img (PIL.Image.Image): The image to place on top.
+        margin (int, optional): The pixel margin from the top and right edges. Defaults to 20.
+
+    Returns:
+        PIL.Image.Image: The combined image.
+    """
+    if base_img is None or overlay_img is None:
+        return base_img
+    
+    base = base_img.convert("RGBA")
+    overlay = overlay_img.convert("RGBA")
+    
+    # --- MODIFICATION ---
+    # Calculate the target width to be 1/5th of the base image's width
+    target_width = base.width // 5
+    
+    # Keep aspect ratio, resize overlay to the newly calculated target width
+    w, h = overlay.size
+    
+    # Add a check to prevent division by zero if the overlay image has no width
+    if w == 0:
+        return base
+        
+    new_height = int(h * (target_width / w))
+    overlay = overlay.resize((target_width, new_height), Image.LANCZOS)
+
+    # Position: top-right corner with a margin
+    x = base.width - overlay.width - margin
+    y = margin
+
+    # Paste the resized overlay onto the base image using its alpha channel for transparency
+    base.paste(overlay, (x, y), overlay)
+    return base
+
+
+# def add_overlay(base_img, overlay_img, margin=20, target_width=200):
+#     if base_img is None or overlay_img is None:
+#         return base_img
+    
+#     base = base_img.convert("RGBA")
+#     overlay = overlay_img.convert("RGBA")
+    
+#     # Keep aspect ratio, resize overlay to target width
+#     w, h = overlay.size
+#     new_height = int(h * (target_width / w))
+#     overlay = overlay.resize((target_width, new_height), Image.LANCZOS)
+
+#     # Position: top-right with margin
+#     x = base.width - overlay.width - margin
+#     y = margin
+
+#     # Paste overlay on base with transparency
+#     base.paste(overlay, (x, y), overlay)
+#     return base
+
+
+# @spaces.GPU
+# def infer(input_image, input_image_upload, overlay_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, 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_upload is not None:
+#         input_image_upload = input_image
+#     elif "composite" in input_image and input_image["composite"] is not None:
+#         input_image = input_image["composite"]
+#     elif "background" in input_image and input_image["background"] is not None:
+#         input_image = input_image["background"]
+#     else:
+#         raise ValueError("No valid image found in EditorValue dict (both 'composite' and 'background' are None)")
+    
+    
+#     if input_image is not None:
+#         if overlay_image is not None:
+#             input_image = add_overlay(input_image, overlay_image)
+            
+#         input_image = input_image.convert("RGB")
+#         image = pipe(
+#             image=input_image, 
+#             prompt=prompt,
+#             guidance_scale=guidance_scale,
+#             width = input_image.size[0],
+#             height = input_image.size[1],
+#             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]
+#     return image, input_image, seed, gr.Button(visible=True)
+
+
+@spaces.GPU
+def infer(input_image, input_image_upload, overlay_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, 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 (dict or PIL.Image.Image): The input from the gr.Paint component.
+        input_image_upload (PIL.Image.Image): The input from the gr.Image upload component.
+        overlay_image (PIL.Image.Image): The face photo to overlay.
+        prompt (str): Text description of the desired edit to apply to the image.
+        seed (int, optional): Random seed for reproducible generation.
+        randomize_seed (bool, optional): If True, generates a random seed.
+        guidance_scale (float, optional): Controls how closely the model follows the prompt.
+        steps (int, optional): Controls how many steps to run the diffusion model for.
+        progress (gr.Progress, optional): Gradio progress tracker.
+    
+    Returns:
+        tuple: A 4-tuple containing the result image, the processed input image, the seed, and a gr.Button update.
+    """
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    # --- CORRECTED LOGIC STARTS HERE ---
+    
+    # 1. Prioritize the uploaded image. If it exists, it becomes our main 'input_image'.
+    if input_image_upload is not None:
+        processed_input_image = input_image_upload
+    # 2. If no image was uploaded, check the drawing canvas.
+    elif isinstance(input_image, dict):
+        # Extract the actual image from the dictionary provided by gr.Paint
+        if "composite" in input_image and input_image["composite"] is not None:
+            processed_input_image = input_image["composite"]
+        elif "background" in input_image and input_image["background"] is not None:
+            processed_input_image = input_image["background"]
+        else:
+            # The canvas is empty, so there's no input image.
+            processed_input_image = None
+    else:
+        # Fallback in case the input is neither from upload nor a valid canvas dict.
+        processed_input_image = None
+        
+    # --- CORRECTED LOGIC ENDS HERE ---
+    
+    # From this point on, 'processed_input_image' is either a PIL Image or None.
+    if processed_input_image is not None:
+        if overlay_image is not None:
+            # Now this function is guaranteed to receive a PIL Image.
+            processed_input_image = add_overlay(processed_input_image, overlay_image)
+            
+        processed_input_image = processed_input_image.convert("RGB")
+        image = pipe(
+            image=processed_input_image, 
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            width = processed_input_image.size[0],
+            height = processed_input_image.size[1],
+            num_inference_steps=steps,
+            generator=torch.Generator().manual_seed(seed),
+        ).images[0]
+    else:
+        # Handle the text-to-image case where no input image was provided.
+        image = pipe(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=steps,
+            generator=torch.Generator().manual_seed(seed),
+        ).images[0]
+        
+    return image, processed_input_image, seed, gr.Button(visible=True)
+    
+@spaces.GPU
+def infer_example(input_image, prompt):
+    image, seed, _ = infer(input_image, prompt)
+    return image, seed
+
+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():
+                    with gr.Tabs() as tabs:
+                        with gr.TabItem("Draw"):
+                            input_image = gr.Paint(
+                                    type="pil",
+                                    brush=gr.Brush(default_size=6, colors=["#000000"], color_mode="fixed"),
+                                    canvas_size = (1200,1200),
+                                    layers = False
+                                )
+                        with gr.TabItem("Upload"):
+                            input_image_upload = gr.Image(label="Upload the drawing", type="pil")
+                with gr.Row():
+                    overlay_image = gr.Image(label="Upload face photo", type="pil")
+                with gr.Row():
+                    prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=1,
+                        value = "make it real",
+                        placeholder="Enter your prompt for editing (e.g., 'Remove glasses', 'Add a hat')",
+                        container=False,
+                    )
+                    run_button = gr.Button("Run", 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=28,
+                        step=1
+                    )
+                    
+            with gr.Column():
+                result = gr.Image(label="Result", show_label=False, interactive=False)
+                result_input = gr.Image(label="Result", show_label=False, interactive=False)
+                reuse_button = gr.Button("Reuse this image", visible=False)
+        
+            
+        examples = gr.Examples(
+            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, seed],
+            fn=infer_example,
+            cache_examples="lazy"
+        )
+            
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn = infer,
+        inputs = [input_image, input_image_upload, overlay_image, prompt, seed, randomize_seed, guidance_scale, steps],
+        outputs = [result, result_input, seed, reuse_button]
+    )
+    reuse_button.click(
+        fn = lambda image: image,
+        inputs = [result],
+        outputs = [input_image]
+    )
+
+demo.launch(mcp_server=True)

optimization.py

@@ -0,0 +1,60 @@
+"""
+"""
+
+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

@@ -0,0 +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 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

@@ -0,0 +1,6 @@
+transformers
+git+https://github.com/huggingface/diffusers.git
+accelerate
+safetensors
+sentencepiece
+peft