Upload 5 files

app.py

@@ -0,0 +1,142 @@
+# PyTorch 2.8 and dependencies (temporary hack)
+import os
+os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces peft')
+
+# Actual demo code
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import random
+from PIL import Image
+
+from diffusers import FluxKontextPipeline
+from diffusers.utils import load_image
+
+from optimization import optimize_pipeline_
+
+MAX_SEED = np.iinfo(np.int32).max
+
+# 1. 加载基础模型
+pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
+
+# 2. 加载 LoRA
+try:
+    pipe.load_lora_weights(".", weight_name="change_clothes_to_nothing_000012800.safetensors")
+    print("Successfully loaded LoRA weights from the root directory.")
+except Exception as e:
+    print(f"Could not load LoRA weights. Please ensure 'change_clothes_to_nothing_000012800.safetensors' is in the root directory. Error: {e}")
+
+
+# 3. 对加载了 LoRA 的模型进行优化
+optimize_pipeline_(pipe, image=Image.new("RGB", (512, 512)), prompt='prompt')
+
+@spaces.GPU
+def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, lora_scale=1.0, progress=gr.Progress(track_tqdm=True)):
+    """
+    使用 FLUX.1 Kontext pipeline 执行图像编辑。
+    """
+    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],
+            height = input_image.size[1],
+            num_inference_steps=steps,
+            generator=torch.Generator().manual_seed(seed),
+            cross_attention_kwargs={"scale": lora_scale}, # 应用 LoRA 强度
+        ).images[0]
+    else:
+        image = pipe(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=steps,
+            generator=torch.Generator().manual_seed(seed),
+            cross_attention_kwargs={"scale": lora_scale}, # 应用 LoRA 强度
+        ).images[0]
+    return image, seed, gr.Button(visible=True)
+
+
+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="上传要编辑的图片", type="pil")
+                with gr.Row():
+                    prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=1,
+                        placeholder="输入您的编辑指令 (例如: '移除眼镜', '添加一顶帽子')",
+                        container=False,
+                    )
+                    run_button = gr.Button("运行", scale=0)
+                with gr.Accordion("高级设置", open=False):
+                    
+                    lora_scale = gr.Slider(
+                        label="LoRA 强度 (LoRA Scale)",
+                        minimum=0.0,
+                        maximum=2.0,
+                        step=0.05,
+                        value=0.8,
+                    )
+
+                    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="结果", show_label=False, interactive=False)
+                reuse_button = gr.Button("复用此图", visible=False)
+            
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn = infer,
+        inputs = [input_image, prompt, seed, randomize_seed, guidance_scale, steps, lora_scale],
+        outputs = [result, seed, reuse_button]
+    )
+    reuse_button.click(
+        fn = lambda image: image,
+        inputs = [result],
+        outputs = [input_image]
+    )
+
+demo.launch(mcp_server=True)

change_clothes_to_nothing_000012800.safetensors

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

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