inb version init

README.md

@@ -1,14 +1,13 @@
 ---
 title: KV Edit
-emoji: 💻
+emoji: 🐢
 colorFrom: gray
-colorTo: red
+colorTo: pink
 sdk: gradio
-sdk_version: 5.17.0
+sdk_version: 5.16.2
 app_file: app.py
 pinned: false
 license: apache-2.0
-short_description: inversion base version
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,7 +1,252 @@
+import os
+import re
+import time
+from dataclasses import dataclass
+from glob import iglob
+from einops import rearrange
+from PIL import ExifTags, Image
+import torch
 import gradio as gr
+import numpy as np
+from flux.sampling import prepare
+from flux.util import (load_ae, load_clip, load_t5)
+from models.kv_edit import Flux_kv_edit,Flux_kv_edit_inf
+import spaces
+from huggingface_hub import login
+login(token=os.getenv('Token'))
 
-def greet(name):
-    return "Hello " + name + "!!"
+@dataclass
+class SamplingOptions:
+    source_prompt: str = ''
+    target_prompt: str = ''
+    # prompt: str
+    width: int = 1366
+    height: int = 768
+    inversion_num_steps: int = 0
+    denoise_num_steps: int = 0
+    skip_step: int = 0
+    inversion_guidance: float = 1.0
+    denoise_guidance: float = 1.0
+    seed: int = 42
+    re_init: bool = False
+    attn_mask: bool = False
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+@torch.inference_mode()
+def encode(init_image, torch_device):
+    init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 127.5 - 1
+    init_image = init_image.unsqueeze(0) 
+    init_image = init_image.to(torch_device)
+    with torch.no_grad():
+        init_image = ae.encode(init_image.to()).to(torch.bfloat16)
+    return init_image
+
+# init all components
+device = "cuda" if torch.cuda.is_available() else "cpu"
+name = 'flux-dev'
+ae = load_ae(name, device)
+t5 = load_t5(device, max_length=256 if name == "flux-schnell" else 512)
+clip = load_clip(device)
+model = Flux_kv_edit(device=device, name=name)
+offload = False
+name = "flux-dev"
+is_schnell = False
+feature_path = 'feature'
+output_dir = 'result'
+add_sampling_metadata = True
+    
+@spaces.GPU(duration=120)
+@torch.inference_mode()
+def edit(init_image, brush_canvas,
+            source_prompt, target_prompt, 
+            inversion_num_steps, denoise_num_steps, 
+            skip_step, 
+            inversion_guidance, denoise_guidance,seed,
+            re_init,attn_mask
+            ):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.cuda.empty_cache()
+    
+    shape = init_image.shape        
+    height = shape[0] if shape[0] % 16 == 0 else shape[0] - shape[0] % 16
+    width = shape[1] if shape[1] % 16 == 0 else shape[1] - shape[1] % 16
+
+    init_image = init_image[:height, :width, :]
+    brush_canvas = brush_canvas["composite"][:,:,:3][:height, :width, :]
+    # 如果brush_Canvas是三通道黑白图，说明就是输入的mask
+    
+    if np.all(brush_canvas[:,:,0] == brush_canvas[:,:,1]) and np.all(brush_canvas[:,:,1] == brush_canvas[:,:,2]):
+        mask = brush_canvas[:,:,0]/255
+        mask = mask.astype(int)
+    else:
+        mask = np.any(init_image != brush_canvas, axis=-1)  # 得到一个二维的布尔数组
+        mask = mask.astype(int)
+    mask_array = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
+    mask_array[:,:,0] = mask * 255  # R
+    mask_array[:,:,3] = mask * 128  # A (半透明，128表示50%透明度)
+    mask_image = Image.fromarray(mask_array, 'RGBA')     
+    original_image = Image.fromarray(np.concatenate((init_image, np.full((height, width, 1), 255, dtype=np.uint8)), axis=2), 'RGBA')
+    masked_image = Image.alpha_composite(original_image, mask_image)
+    mask = torch.from_numpy(mask).unsqueeze(0).unsqueeze(0).to(torch.bfloat16).to(device)
+    
+    init_image = encode(init_image, device).to(device)
+    
+    seed = int(seed)
+    if seed == -1:
+        seed = torch.randint(0, 2**32, (1,)).item()
+    opts = SamplingOptions(
+        source_prompt=source_prompt,
+        target_prompt=target_prompt,
+        width=width,
+        height=height,
+        inversion_num_steps=inversion_num_steps,
+        denoise_num_steps=denoise_num_steps,
+        skip_step=skip_step,
+        inversion_guidance=inversion_guidance,
+        denoise_guidance=denoise_guidance,
+        seed=seed,
+        re_init=re_init,
+        attn_mask=attn_mask
+    )
+
+        
+    torch.manual_seed(opts.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(opts.seed)
+
+    t0 = time.perf_counter()
+
+    #############inverse#######################
+    # 将布尔数组转换为整数类型，如果需要1和0而不是True和False的话
+    with torch.no_grad():
+        inp = prepare(t5, clip, init_image, prompt=opts.source_prompt)
+        inp_target = prepare(t5, clip, init_image, prompt=opts.target_prompt)
+      
+    x = model(inp, inp_target, mask, opts)
+    
+    device = torch.device("cuda")
+    with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+        x = ae.decode(x)
+    # 得到还在显卡上的特征
+    # bring into PIL format and save
+    x = x.clamp(-1, 1)
+    # x = embed_watermark(x.float())
+    x = x.float().cpu()
+    x = rearrange(x[0], "c h w -> h w c")
+    
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    #############回到像素空间就算结束#######################
+    output_name = os.path.join(output_dir, "img_{idx}.jpg")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        idx = 0
+    else:
+        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
+        if len(fns) > 0:
+            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
+        else:
+            idx = 0
+    #############找idx#######################
+    
+    fn = output_name.format(idx=idx)
+
+    img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
+    exif_data = Image.Exif()
+    exif_data[ExifTags.Base.Software] = "AI generated;txt2img;flux"
+    exif_data[ExifTags.Base.Make] = "Black Forest Labs"
+    exif_data[ExifTags.Base.Model] = name
+
+    exif_data[ExifTags.Base.ImageDescription] = source_prompt
+    img.save(fn, exif=exif_data, quality=95, subsampling=0)
+    masked_image.save(fn.replace(".jpg", "_mask.png"),  format='PNG')
+    t1 = time.perf_counter()
+    print(f"Done in {t1 - t0:.1f}s. Saving {fn}")
+    
+    print("End Edit")
+    return img
+
+
+    
+def create_demo(model_name: str):
+    # editor = FluxEditor_kv_demo()
+    is_schnell = model_name == "flux-schnell"
+    
+    title = r"""
+        <h1 align="center">🎨 KV-Edit: Training-Free Image Editing for Precise Background Preservation</h1>
+        """
+        
+    description = r"""
+        <b>Official 🤗 Gradio demo</b> for <a href='https://github.com/Xilluill/KV-Edit' target='_blank'><b>KV-Edit: Training-Free Image Editing for Precise Background Preservation</b></a>.<br>
+    
+        🔔🔔[<b>Important</b>] Editing steps:<br>
+        1️⃣ Upload your image that needs to be edited (The resolution is expected be less than 1360*768, or the memory of GPU may be not enough.) <br>
+        2️⃣ Re-upload the original image and use the brush tool to draw your mask area. <br>
+        3️⃣ Fill in your source prompt and target prompt, then adjust the hyperparameters. <br>
+        4️⃣ Click the "Edit" button to generate your edited image! <br>
+        """
+    article = r"""
+    If our work is helpful, please help to ⭐ the <a href='https://github.com/Xilluill/KV-Edit' target='_blank'>Github Repo</a>. Thanks! 
+    """
+
+    badge = r"""
+    [![GitHub Stars](https://img.shields.io/github/stars/Xilluill/KV-Edit)](https://github.com/Xilluill/KV-Edit)
+    """
+    
+    with gr.Blocks() as demo:
+        gr.HTML(title)
+        gr.Markdown(description)
+        gr.Markdown(article)
+        # gr.Markdown(badge)
+        
+        with gr.Row():
+            with gr.Column():
+                source_prompt = gr.Textbox(label="Source Prompt", value='In a cluttered wooden cabin, a workbench holds a green neon sign that reads "I love here"' )
+                inversion_num_steps = gr.Slider(1, 50, 28, step=1, label="Number of inversion steps")
+                target_prompt = gr.Textbox(label="Target Prompt", value='In a cluttered wooden cabin, a workbench holds a green neon sign that reads "I love iccv"' )
+                denoise_num_steps = gr.Slider(1, 50, 28, step=1, label="Number of denoise steps")
+                init_image = gr.Image(label="Input Image", visible=True)
+                brush_canvas = gr.ImageEditor(label="Brush Canvas",
+                                                sources=('upload'), 
+                                                brush=gr.Brush(default_size=10,
+                                                               default_color="#000000"),
+                                                interactive=True,
+                                                container=True,
+                                                transforms=[],
+                                                height="auto",
+                                                format='png',scale=1)
+                
+                edit_btn = gr.Button("edit")
+                
+                
+            with gr.Column():
+                with gr.Accordion("Advanced Options", open=True):
+                    # num_steps = gr.Slider(1, 30, 25, step=1, label="Number of steps")
+
+                    skip_step = gr.Slider(0, 30, 4, step=1, label="Number of inject steps")
+                    inversion_guidance = gr.Slider(1.0, 10.0, 1.5, step=0.1, label="inversion Guidance", interactive=not is_schnell)
+                    denoise_guidance = gr.Slider(1.0, 10.0, 5.5, step=0.1, label="denoise Guidance", interactive=not is_schnell)
+                    seed = gr.Textbox('0', label="Seed (-1 for random)", visible=True)
+                    with gr.Row():
+                        re_init = gr.Checkbox(label="re_init", value=False)
+                        attn_mask = gr.Checkbox(label="attn_mask", value=False)
+
+                
+                output_image = gr.Image(label="Generated Image")
+        edit_btn.click(
+            fn=edit,
+            inputs=[init_image, brush_canvas,
+                    source_prompt, target_prompt, 
+                    inversion_num_steps, denoise_num_steps, 
+                    skip_step, 
+                    inversion_guidance,
+                    denoise_guidance,seed,
+                    re_init,attn_mask
+                    ],
+            outputs=[output_image]
+        )
+    return demo
+
+
+demo = create_demo("flux-dev")
+
+demo.launch()

flux/__init__.py

@@ -0,0 +1,11 @@
+try:
+    from ._version import version as __version__  # type: ignore
+    from ._version import version_tuple
+except ImportError:
+    __version__ = "unknown (no version information available)"
+    version_tuple = (0, 0, "unknown", "noinfo")
+
+from pathlib import Path
+
+PACKAGE = __package__.replace("_", "-")
+PACKAGE_ROOT = Path(__file__).parent

flux/__main__.py

@@ -0,0 +1,4 @@
+from .cli import app
+
+if __name__ == "__main__":
+    app()

flux/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.0.post6+ge52e00f.d20250111'
+__version_tuple__ = version_tuple = (0, 0, 'ge52e00f.d20250111')

flux/api.py

@@ -0,0 +1,194 @@
+import io
+import os
+import time
+from pathlib import Path
+
+import requests
+from PIL import Image
+
+API_ENDPOINT = "https://api.bfl.ml"
+
+
+class ApiException(Exception):
+    def __init__(self, status_code: int, detail: str | list[dict] | None = None):
+        super().__init__()
+        self.detail = detail
+        self.status_code = status_code
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+    def __repr__(self) -> str:
+        if self.detail is None:
+            message = None
+        elif isinstance(self.detail, str):
+            message = self.detail
+        else:
+            message = "[" + ",".join(d["msg"] for d in self.detail) + "]"
+        return f"ApiException({self.status_code=}, {message=}, detail={self.detail})"
+
+
+class ImageRequest:
+    def __init__(
+        self,
+        prompt: str,
+        width: int = 1024,
+        height: int = 1024,
+        name: str = "flux.1-pro",
+        num_steps: int = 50,
+        prompt_upsampling: bool = False,
+        seed: int | None = None,
+        validate: bool = True,
+        launch: bool = True,
+        api_key: str | None = None,
+    ):
+        """
+        Manages an image generation request to the API.
+
+        Args:
+            prompt: Prompt to sample
+            width: Width of the image in pixel
+            height: Height of the image in pixel
+            name: Name of the model
+            num_steps: Number of network evaluations
+            prompt_upsampling: Use prompt upsampling
+            seed: Fix the generation seed
+            validate: Run input validation
+            launch: Directly launches request
+            api_key: Your API key if not provided by the environment
+
+        Raises:
+            ValueError: For invalid input
+            ApiException: For errors raised from the API
+        """
+        if validate:
+            if name not in ["flux.1-pro"]:
+                raise ValueError(f"Invalid model {name}")
+            elif width % 32 != 0:
+                raise ValueError(f"width must be divisible by 32, got {width}")
+            elif not (256 <= width <= 1440):
+                raise ValueError(f"width must be between 256 and 1440, got {width}")
+            elif height % 32 != 0:
+                raise ValueError(f"height must be divisible by 32, got {height}")
+            elif not (256 <= height <= 1440):
+                raise ValueError(f"height must be between 256 and 1440, got {height}")
+            elif not (1 <= num_steps <= 50):
+                raise ValueError(f"steps must be between 1 and 50, got {num_steps}")
+
+        self.request_json = {
+            "prompt": prompt,
+            "width": width,
+            "height": height,
+            "variant": name,
+            "steps": num_steps,
+            "prompt_upsampling": prompt_upsampling,
+        }
+        if seed is not None:
+            self.request_json["seed"] = seed
+
+        self.request_id: str | None = None
+        self.result: dict | None = None
+        self._image_bytes: bytes | None = None
+        self._url: str | None = None
+        if api_key is None:
+            self.api_key = os.environ.get("BFL_API_KEY")
+        else:
+            self.api_key = api_key
+
+        if launch:
+            self.request()
+
+    def request(self):
+        """
+        Request to generate the image.
+        """
+        if self.request_id is not None:
+            return
+        response = requests.post(
+            f"{API_ENDPOINT}/v1/image",
+            headers={
+                "accept": "application/json",
+                "x-key": self.api_key,
+                "Content-Type": "application/json",
+            },
+            json=self.request_json,
+        )
+        result = response.json()
+        if response.status_code != 200:
+            raise ApiException(status_code=response.status_code, detail=result.get("detail"))
+        self.request_id = response.json()["id"]
+
+    def retrieve(self) -> dict:
+        """
+        Wait for the generation to finish and retrieve response.
+        """
+        if self.request_id is None:
+            self.request()
+        while self.result is None:
+            response = requests.get(
+                f"{API_ENDPOINT}/v1/get_result",
+                headers={
+                    "accept": "application/json",
+                    "x-key": self.api_key,
+                },
+                params={
+                    "id": self.request_id,
+                },
+            )
+            result = response.json()
+            if "status" not in result:
+                raise ApiException(status_code=response.status_code, detail=result.get("detail"))
+            elif result["status"] == "Ready":
+                self.result = result["result"]
+            elif result["status"] == "Pending":
+                time.sleep(0.5)
+            else:
+                raise ApiException(status_code=200, detail=f"API returned status '{result['status']}'")
+        return self.result
+
+    @property
+    def bytes(self) -> bytes:
+        """
+        Generated image as bytes.
+        """
+        if self._image_bytes is None:
+            response = requests.get(self.url)
+            if response.status_code == 200:
+                self._image_bytes = response.content
+            else:
+                raise ApiException(status_code=response.status_code)
+        return self._image_bytes
+
+    @property
+    def url(self) -> str:
+        """
+        Public url to retrieve the image from
+        """
+        if self._url is None:
+            result = self.retrieve()
+            self._url = result["sample"]
+        return self._url
+
+    @property
+    def image(self) -> Image.Image:
+        """
+        Load the image as a PIL Image
+        """
+        return Image.open(io.BytesIO(self.bytes))
+
+    def save(self, path: str):
+        """
+        Save the generated image to a local path
+        """
+        suffix = Path(self.url).suffix
+        if not path.endswith(suffix):
+            path = path + suffix
+        Path(path).resolve().parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "wb") as file:
+            file.write(self.bytes)
+
+
+if __name__ == "__main__":
+    from fire import Fire
+
+    Fire(ImageRequest)

flux/math.py

@@ -0,0 +1,40 @@
+import torch
+from einops import rearrange
+from torch import Tensor
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor,pe_q = None, attention_mask = None) -> Tensor:
+    if pe_q is None:
+        q, k = apply_rope(q, k, pe) # torch.Size([1, 24, 4592, 128])  torch.Size([1, 24, 4592, 128]) pe torch.Size([1, 1, 4592, 64, 2, 2])
+        x = torch.nn.functional.scaled_dot_product_attention(q, k, v,attn_mask=attention_mask) # torch.Size([1, 24, 4592, 128])
+        x = rearrange(x, "B H L D -> B L (H D)") # torch.Size([1, 4592, 3072])
+        return x
+    else: 
+        q, k = apply_rope_qk(q, k, pe_q, pe) # torch.Size([1, 24, 4592, 128])  torch.Size([1, 24, 4592, 128]) pe torch.Size([1, 1, 4592, 64, 2, 2])
+        x = torch.nn.functional.scaled_dot_product_attention(q, k, v,attn_mask=attention_mask)
+        x = rearrange(x, "B H L D -> B L (H D)")
+        return x
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim # dim =16 + 56 + 56 
+    omega = 1.0 / (theta**scale) # 64 omega
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2) # torch.Size([1, 1, 4592, x, 2, 2]) x = 8 + 28 + 28
+    return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) # torch.Size([1, 24, 4592, 128]) -> torch.Size([1, 24, 4592, 64, 1, 2])
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) # torch.Size([1, 24, 4592, 128]) -> torch.Size([1, 24, 4592, 64, 1, 2])
+    xq_out = freqs_cis[:, :, :xq_.shape[2], :, :, 0] * xq_[..., 0] + freqs_cis[:, :, :xq_.shape[2], :, :, 1] * xq_[..., 1] 
+    xk_out = freqs_cis[:, :, :xk_.shape[2], :, :, 0] * xk_[..., 0] + freqs_cis[:, :, :xk_.shape[2], :, :, 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+def apply_rope_qk(xq: Tensor, xk: Tensor, freqs_cis_q: Tensor,freqs_cis_k: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) # torch.Size([1, 24, 4592, 128]) -> torch.Size([1, 24, 4592, 64, 1, 2])
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) # torch.Size([1, 24, 4592, 128]) -> torch.Size([1, 24, 4592, 64, 1, 2])
+    xq_out = freqs_cis_q[:, :, :xq_.shape[2], :, :, 0] * xq_[..., 0] + freqs_cis_q[:, :, :xq_.shape[2], :, :, 1] * xq_[..., 1]  
+    xk_out = freqs_cis_k[:, :, :xk_.shape[2], :, :, 0] * xk_[..., 0] + freqs_cis_k[:, :, :xk_.shape[2], :, :, 1] * xk_[..., 1] 
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)

flux/model.py

@@ -0,0 +1,170 @@
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+
+from flux.modules.layers import (DoubleStreamBlock, EmbedND, LastLayer,
+                                 MLPEmbedder, SingleStreamBlock,DoubleStreamBlock_rf,SingleStreamBlock_rf,
+                                    SingleStreamBlock_kv,DoubleStreamBlock_kv,
+                                 timestep_embedding)
+
+
+@dataclass
+class FluxParams:
+    in_channels: int
+    vec_in_dim: int
+    context_in_dim: int
+    hidden_size: int
+    mlp_ratio: float
+    num_heads: int
+    depth: int
+    depth_single_blocks: int
+    axes_dim: list[int]
+    theta: int
+    qkv_bias: bool
+    guidance_embed: bool
+
+
+class Flux(nn.Module):
+    """
+    Transformer model for flow matching on sequences.
+    """
+
+    def __init__(self, params: FluxParams,double_block_cls=DoubleStreamBlock,single_block_cls=SingleStreamBlock):
+        super().__init__()
+
+        self.params = params
+        self.in_channels = params.in_channels
+        self.out_channels = self.in_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(
+                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+            )
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
+        self.guidance_in = (
+            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
+        )
+        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
+
+        self.double_blocks = nn.ModuleList(
+            [
+                double_block_cls(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.single_blocks = nn.ModuleList(
+            [
+                single_block_cls(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
+                for _ in range(params.depth_single_blocks)
+            ]
+        )
+
+        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
+
+    def forward(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+        y: Tensor,
+        guidance: Tensor | None = None,
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        for block in self.double_blocks:
+            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+
+        img = torch.cat((txt, img), 1)
+        for block in self.single_blocks:
+            img = block(img, vec=vec, pe=pe)
+        img = img[:, txt.shape[1] :, ...]
+
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+        return img
+
+class Flux_kv(Flux):
+    """
+    继承Flux类，重写forward方法
+    """
+
+    def __init__(self, params: FluxParams,double_block_cls=DoubleStreamBlock_kv,single_block_cls=SingleStreamBlock_kv):
+        super().__init__(params,double_block_cls,single_block_cls)
+
+    def forward(
+        self,
+        img: Tensor, # （B,x,x） (1,4080,64)
+        img_ids: Tensor,
+        txt: Tensor, # torch.Size([1, 512, 4096])
+        txt_ids: Tensor,
+        timesteps: Tensor, # torch.Size([1])
+        y: Tensor, # torch.Size([1, 768])
+        guidance: Tensor | None = None, # torch.Size([1])
+        info: dict = {},
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+        vec = self.time_in(timestep_embedding(timesteps, 256)) # torch.Size([1, 3072])
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256)) # torch.Size([1, 3072])
+        vec = vec + self.vector_in(y)# torch.Size([1, 3072])
+        txt = self.txt_in(txt) # ([1, 512, 4096]) -> torch.Size([1, 512, 3072])
+
+        ids = torch.cat((txt_ids, img_ids), dim=1) # torch.Size([1, 512, 3])  torch.Size([1, 4080, 3]) -> torch.Size([1, 4592, 3])
+        pe = self.pe_embedder(ids) # torch.Size([1, 1, 4592, 64, 2, 2])
+        if not info['inverse']:
+            mask_indices = info['mask_indices'] # 图片seq坐标下的
+            info['pe_mask'] = torch.cat((pe[:, :, :512, ...],pe[:, :, mask_indices+512, ...]),dim=2)
+
+        cnt = 0
+        for block in self.double_blocks:
+            info['id'] = cnt
+            img, txt = block(img=img, txt=txt, vec=vec, pe=pe, info=info)
+            cnt += 1
+
+        cnt = 0
+        x = torch.cat((txt, img), 1) 
+        for block in self.single_blocks:
+            info['id'] = cnt
+            x = block(x, vec=vec, pe=pe, info=info)
+            cnt += 1
+
+        img = x[:, txt.shape[1] :, ...]
+
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+
+        return img

flux/modules/autoencoder.py

@@ -0,0 +1,313 @@
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+
+@dataclass
+class AutoEncoderParams:
+    resolution: int
+    in_channels: int
+    ch: int
+    out_ch: int
+    ch_mult: list[int]
+    num_res_blocks: int
+    z_channels: int
+    scale_factor: float
+    shift_factor: float
+
+
+def swish(x: Tensor) -> Tensor:
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = swish(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = swish(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        # no asymmetric padding in torch conv, must do it ourselves
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor):
+        pad = (0, 1, 0, 1)
+        x = nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        # downsampling
+        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class DiagonalGaussian(nn.Module):
+    def __init__(self, sample: bool = True, chunk_dim: int = 1):
+        super().__init__()
+        self.sample = sample
+        self.chunk_dim = chunk_dim
+
+    def forward(self, z: Tensor) -> Tensor:
+        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
+        # import pdb;pdb.set_trace()
+        if self.sample:
+            std = torch.exp(0.5 * logvar)
+            return mean #+ std * torch.randn_like(mean)
+        else:
+            return mean
+
+
+class AutoEncoder(nn.Module):
+    def __init__(self, params: AutoEncoderParams):
+        super().__init__()
+        self.encoder = Encoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.decoder = Decoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            out_ch=params.out_ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.reg = DiagonalGaussian()
+
+        self.scale_factor = params.scale_factor
+        self.shift_factor = params.shift_factor
+
+    def encode(self, x: Tensor) -> Tensor:
+        z = self.reg(self.encoder(x))
+        z = self.scale_factor * (z - self.shift_factor)
+        return z
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = z / self.scale_factor + self.shift_factor
+        return self.decoder(z)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.decode(self.encode(x))

flux/modules/conditioner.py

@@ -0,0 +1,38 @@
+from torch import Tensor, nn
+from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel,
+                          T5Tokenizer)
+
+
+class HFEmbedder(nn.Module):
+    def __init__(self, version: str, max_length: int, is_clip, **hf_kwargs):
+        super().__init__()
+        self.is_clip = is_clip
+        self.max_length = max_length
+        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
+
+        if self.is_clip:
+            self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
+            self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
+        else:
+            self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
+            self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
+
+        self.hf_module = self.hf_module.eval().requires_grad_(False)
+
+    def forward(self, text: list[str]) -> Tensor:
+        batch_encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            return_length=False,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+        outputs = self.hf_module(
+            input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
+            attention_mask=None,
+            output_hidden_states=False,
+        )
+        return outputs[self.output_key]

flux/modules/layers.py

@@ -0,0 +1,386 @@
+import math
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+from flux.math import attention, rope,apply_rope
+
+import os
+
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], # [1, 1, 4592, 8, 2, 2]) [1, 1, 4592, 28, 2, 2]) [1, 1, 4592, 28, 2, 2])
+            dim=-3,
+        )
+
+        return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+    """
+    Create sinusoidal timestep embeddings.
+    :param t: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an (N, D) Tensor of positional embeddings.
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
+        t.device
+    )
+
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor):
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.norm = QKNorm(head_dim)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
+        qkv = self.qkv(x)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+        x = attention(q, k, v, pe=pe)
+        x = self.proj(x)
+        return x
+
+
+@dataclass
+class ModulationOut:
+    shift: Tensor
+    scale: Tensor
+    gate: Tensor
+
+
+class Modulation(nn.Module):
+    def __init__(self, dim: int, double: bool):
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
+        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
+
+        return (
+            ModulationOut(*out[:3]),
+            ModulationOut(*out[3:]) if self.is_double else None,
+        )
+
+
+class DoubleStreamBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
+        super().__init__()
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.img_mod = Modulation(hidden_size, double=True)
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+        self.txt_mod = Modulation(hidden_size, double=True)
+        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        img_mod1, img_mod2 = self.img_mod(vec)
+        txt_mod1, txt_mod2 = self.txt_mod(vec)
+
+        # prepare image for attention
+        img_modulated = self.img_norm1(img)
+        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) # torch.Size([1, 24, 4080, 128])
+
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+        # prepare txt for attention
+        txt_modulated = self.txt_norm1(txt)
+        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        # run actual attention
+        q = torch.cat((txt_q, img_q), dim=2) # [8, 24, 512, 128] + [8, 24, 900, 128] -> [8, 24, 1412, 128]
+        k = torch.cat((txt_k, img_k), dim=2)
+        v = torch.cat((txt_v, img_v), dim=2)
+        # import pdb;pdb.set_trace()
+        attn = attention(q, k, v, pe=pe)
+ 
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        # calculate the img bloks
+        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+
+        # calculate the txt bloks
+        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+        return img, txt
+class SingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers as described in
+    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qk_scale: float | None = None,
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        # qkv and mlp_in
+        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
+        # proj and mlp_out
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
+
+        self.norm = QKNorm(head_dim)
+
+        self.hidden_size = hidden_size
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.mlp_act = nn.GELU(approximate="tanh")
+        self.modulation = Modulation(hidden_size, double=False)
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        mod, _ = self.modulation(vec)
+        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+
+        # compute attention
+        attn = attention(q, k, v, pe=pe)
+        # compute activation in mlp stream, cat again and run second linear layer
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        return x + mod.gate * output
+
+
+class LastLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+        x = self.linear(x)
+        return x
+
+
+class LastLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+        x = self.linear(x)
+        return x
+
+
+class DoubleStreamBlock_kv(DoubleStreamBlock):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
+        super().__init__(hidden_size, num_heads, mlp_ratio, qkv_bias)
+
+    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, info) -> tuple[Tensor, Tensor]:
+        img_mod1, img_mod2 = self.img_mod(vec)
+        txt_mod1, txt_mod2 = self.txt_mod(vec)
+
+        # prepare image for attention
+        img_modulated = self.img_norm1(img)
+        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) # torch.Size([1, 24, 4080, 128])
+
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+        # prepare txt for attention
+        txt_modulated = self.txt_norm1(txt)
+        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)   
+        
+        feature_k_name = str(info['t']) + '_' + str(info['id']) + '_' + 'MB' + '_' + 'K'
+        feature_v_name = str(info['t']) + '_' + str(info['id']) + '_' + 'MB' + '_' + 'V'
+        if info['inverse']:
+            info['feature'][feature_k_name] = img_k.cpu()
+            info['feature'][feature_v_name] = img_v.cpu()
+            q = torch.cat((txt_q, img_q), dim=2) # [B, 24, 512, 128] + [B, 24, 900, 128] -> [B, 24, 1412, 128]
+            k = torch.cat((txt_k, img_k), dim=2)
+            v = torch.cat((txt_v, img_v), dim=2)
+            if 'attention_mask' in info:
+                attn = attention(q, k, v, pe=pe,attention_mask=info['attention_mask'])
+            else:
+                attn = attention(q, k, v, pe=pe)
+    
+        # elif feature_k_name in info['feature']:
+        else:
+            source_img_k = info['feature'][feature_k_name].to(img.device)
+            source_img_v = info['feature'][feature_v_name].to(img.device)
+    
+            mask_indices = info['mask_indices'] # 图片seq坐标下的
+            source_img_k[:, :, mask_indices, ...] = img_k
+            source_img_v[:, :, mask_indices, ...] = img_v
+            
+            q = torch.cat((txt_q, img_q), dim=2)
+            k = torch.cat((txt_k, source_img_k), dim=2)
+            v = torch.cat((txt_v, source_img_v), dim=2)
+            attn = attention(q, k, v, pe=pe, pe_q = info['pe_mask'])
+
+        
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        # calculate the img bloks
+        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+
+        # calculate the txt bloks
+        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+        return img, txt
+    
+class SingleStreamBlock_kv(SingleStreamBlock):
+    """
+    A DiT block with parallel linear layers as described in
+    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qk_scale: float | None = None,
+    ):
+        super().__init__(hidden_size, num_heads, mlp_ratio, qk_scale)
+
+    def forward(self,x: Tensor, vec: Tensor, pe: Tensor, info) -> Tensor:
+        mod, _ = self.modulation(vec)
+        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+        img_k = k[:, :, 512:, ...]
+        img_v = v[:, :, 512:, ...]
+        
+        txt_k = k[:, :, :512, ...]
+        txt_v = v[:, :, :512, ...]
+    
+
+        feature_k_name = str(info['t']) + '_' + str(info['id']) + '_' + 'SB' + '_' + 'K'
+        feature_v_name = str(info['t']) + '_' + str(info['id']) + '_' + 'SB' + '_' + 'V'
+        if info['inverse']:
+            info['feature'][feature_k_name] = img_k.cpu()
+            info['feature'][feature_v_name] = img_v.cpu()
+            if 'attention_mask' in info:
+                attn = attention(q, k, v, pe=pe,attention_mask=info['attention_mask'])
+            else:
+                attn = attention(q, k, v, pe=pe)
+            
+        else:
+            source_img_k = info['feature'][feature_k_name].to(x.device)
+            source_img_v = info['feature'][feature_v_name].to(x.device)
+        
+            mask_indices = info['mask_indices'] # 图片seq坐标下的
+            source_img_k[:, :, mask_indices, ...] = img_k
+            source_img_v[:, :, mask_indices, ...] = img_v
+            
+            k = torch.cat((txt_k, source_img_k), dim=2)
+            v = torch.cat((txt_v, source_img_v), dim=2)
+            attn = attention(q, k, v, pe=pe, pe_q = info['pe_mask'])
+
+        # compute attention
+        # attn = attention(q, k, v, pe=pe)
+        # compute activation in mlp stream, cat again and run second linear layer
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        return x + mod.gate * output

flux/sampling.py

@@ -0,0 +1,306 @@
+import math
+from typing import Callable
+
+import torch
+from einops import rearrange, repeat
+from torch import Tensor
+
+from .model import Flux,Flux_kv
+from .modules.conditioner import HFEmbedder
+from tqdm import tqdm
+from tqdm.contrib import tzip
+
+def get_noise(
+    num_samples: int,
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+):
+    return torch.randn(
+        num_samples,
+        16,
+        # allow for packing
+        2 * math.ceil(height / 16),
+        2 * math.ceil(width / 16),
+        device=device,
+        dtype=dtype,
+        generator=torch.Generator(device=device).manual_seed(seed),
+    )
+    
+
+def prepare(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str | list[str]) -> dict[str, Tensor]:
+    bs, c, h, w = img.shape
+    if bs == 1 and not isinstance(prompt, str):
+        bs = len(prompt)
+
+    img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    if img.shape[0] == 1 and bs > 1:
+        img = repeat(img, "1 ... -> bs ...", bs=bs)
+
+    img_ids = torch.zeros(h // 2, w // 2, 3)
+    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
+    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+    if isinstance(prompt, str):
+        prompt = [prompt]
+    txt = t5(prompt)
+    if txt.shape[0] == 1 and bs > 1:
+        txt = repeat(txt, "1 ... -> bs ...", bs=bs)
+    txt_ids = torch.zeros(bs, txt.shape[1], 3)
+
+    vec = clip(prompt)
+    if vec.shape[0] == 1 and bs > 1:
+        vec = repeat(vec, "1 ... -> bs ...", bs=bs)
+
+    return {
+        "img": img,
+        "img_ids": img_ids.to(img.device),
+        "txt": txt.to(img.device),
+        "txt_ids": txt_ids.to(img.device),
+        "vec": vec.to(img.device),
+    }
+
+def prepare_flowedit(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, source_prompt: str | list[str],target_prompt) -> dict[str, Tensor]:
+    bs, c, h, w = img.shape
+    if bs == 1 and not isinstance(source_prompt, str):
+        bs = len(source_prompt)
+
+    img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    if img.shape[0] == 1 and bs > 1:
+        img = repeat(img, "1 ... -> bs ...", bs=bs)
+
+    img_ids = torch.zeros(h // 2, w // 2, 3)
+    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
+    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+    # if isinstance(prompt, str):
+    #     prompt = [prompt]
+    # txt = t5(prompt)
+    # if txt.shape[0] == 1 and bs > 1:
+    #     txt = repeat(txt, "1 ... -> bs ...", bs=bs)
+    # txt_ids = torch.zeros(bs, txt.shape[1], 3)
+
+    # vec = clip(prompt)
+    # if vec.shape[0] == 1 and bs > 1:
+    #     vec = repeat(vec, "1 ... -> bs ...", bs=bs)
+    if isinstance(source_prompt, str):
+        source_prompt = [source_prompt]
+    source_txt = t5(source_prompt)
+    if source_txt.shape[0] == 1 and bs > 1:
+        source_txt = repeat(source_txt, "1 ... -> bs ...", bs=bs)
+    source_txt_ids = torch.zeros(bs, source_txt.shape[1], 3)
+    
+    source_vec = clip(target_prompt)
+    if source_vec.shape[0] == 1 and bs > 1:
+        source_vec = repeat(source_vec, "1 ... -> bs ...", bs=bs)
+    
+    if isinstance(target_prompt, str):
+        target_prompt = [target_prompt]
+    target_txt = t5(target_prompt)
+    if target_txt.shape[0] == 1 and bs > 1:
+        target_txt = repeat(target_txt, "1 ... -> bs ...", bs=bs)
+    target_txt_ids = torch.zeros(bs, target_txt.shape[1], 3)
+    
+    target_vec = clip(target_prompt)
+    if target_vec.shape[0] == 1 and bs > 1:
+        target_vec = repeat(target_vec, "1 ... -> bs ...", bs=bs)
+    
+
+    return {
+        "img": img,
+        "img_ids": img_ids.to(img.device),
+        "source_txt": source_txt.to(img.device),
+        "source_txt_ids": source_txt_ids.to(img.device),
+        "source_vec": source_vec.to(img.device),
+        "target_txt": target_txt.to(img.device),
+        "target_txt_ids": target_txt_ids.to(img.device),
+        "target_vec": target_vec.to(img.device)
+    }
+
+def time_shift(mu: float, sigma: float, t: Tensor):
+    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+
+def get_lin_function(
+    x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
+) -> Callable[[float], float]:
+    m = (y2 - y1) / (x2 - x1)
+    b = y1 - m * x1
+    return lambda x: m * x + b
+
+
+def get_schedule(
+    num_steps: int,
+    image_seq_len: int,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+    shift: bool = True,
+) -> list[float]:
+    # extra step for zero
+    timesteps = torch.linspace(1, 0, num_steps + 1)
+
+    # shifting the schedule to favor high timesteps for higher signal images
+    if shift:
+        # estimate mu based on linear estimation between two points
+        mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
+        timesteps = time_shift(mu, 1.0, timesteps)
+
+    return timesteps.tolist()
+
+
+def denoise(
+    model: Flux,
+    # model input
+    img: Tensor,
+    img_ids: Tensor,
+    txt: Tensor,
+    txt_ids: Tensor,
+    vec: Tensor,
+    # sampling parameters
+    timesteps: list[float],
+    guidance: float = 4.0,
+):
+    # this is ignored for schnell
+    guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
+    for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])):
+        t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+        pred = model(
+            img=img,
+            img_ids=img_ids,
+            txt=txt,
+            txt_ids=txt_ids,
+            y=vec,
+            timesteps=t_vec,
+            guidance=guidance_vec,
+        )
+
+        img = img + (t_prev - t_curr) * pred
+
+    return img
+
+def unpack(x: Tensor, height: int, width: int) -> Tensor:
+    return rearrange(
+        x,
+        "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+        h=math.ceil(height / 16),
+        w=math.ceil(width / 16),
+        ph=2,
+        pw=2,
+    )
+
+def denoise_kv(
+    model: Flux_kv,
+    # model input
+    img: Tensor,
+    img_ids: Tensor,
+    txt: Tensor,
+    txt_ids: Tensor,
+    vec: Tensor,
+    # sampling parameters
+    timesteps: list[float],
+    inverse,
+    info, 
+    guidance: float = 4.0
+):
+
+    if inverse:
+        timesteps = timesteps[::-1]
+        
+    guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
+    
+    for i, (t_curr, t_prev) in enumerate(tzip(timesteps[:-1], timesteps[1:])):
+        t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+        info['t'] = t_prev if inverse else t_curr
+        
+        if inverse:
+            img_name = str(info['t']) + '_' + 'img'
+            info['feature'][img_name] = img.cpu()
+        else:
+            img_name = str(info['t']) + '_' + 'img'
+            source_img = info['feature'][img_name].to(img.device)
+            img = source_img[:, info['mask_indices'],...] * (1 - info['mask'][:, info['mask_indices'],...]) + img * info['mask'][:, info['mask_indices'],...]
+        pred = model(
+            img=img,
+            img_ids=img_ids,
+            txt=txt,
+            txt_ids=txt_ids,
+            y=vec,
+            timesteps=t_vec,
+            guidance=guidance_vec,
+            info=info
+        )
+        img = img + (t_prev - t_curr) * pred
+    return img, info
+
+def denoise_kv_inf(
+    model: Flux_kv,
+    # model input
+    img: Tensor,
+    img_ids: Tensor,
+    source_txt: Tensor,
+    source_txt_ids: Tensor,
+    source_vec: Tensor,
+    target_txt: Tensor,
+    target_txt_ids: Tensor,
+    target_vec: Tensor,
+    # sampling parameters
+    timesteps: list[float],
+    target_guidance: float = 4.0,
+    source_guidance: float = 4.0,
+    info: dict = {},
+):
+        
+    target_guidance_vec = torch.full((img.shape[0],), target_guidance, device=img.device, dtype=img.dtype)
+    source_guidance_vec = torch.full((img.shape[0],), source_guidance, device=img.device, dtype=img.dtype)
+    
+    mask_indices = info['mask_indices']
+    init_img = img.clone() # torch.Size([1, 4080, 64])
+    z_fe = img[:, mask_indices,...]
+    
+    noise_list = []
+    for i in range(len(timesteps)):
+        noise = torch.randn(init_img.size(), dtype=init_img.dtype, 
+                        layout=init_img.layout, device=init_img.device,
+                        generator=torch.Generator(device=init_img.device).manual_seed(0)) # 每次重新取噪声 根据t进行加噪
+        noise_list.append(noise)
+
+    for i, (t_curr, t_prev) in enumerate(tzip(timesteps[:-1], timesteps[1:])): # 从高到低
+        
+        info['t'] = 'inf'
+        t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+        
+        z_src = (1 - t_curr) * init_img + t_curr * noise_list[i]
+        z_tar = z_src[:, mask_indices,...] - init_img[:, mask_indices,...] + z_fe
+        
+        info['inverse'] = True
+        info['feature'] = {} # 清空kv特征
+        v_src = model(
+            img=z_src,
+            img_ids=img_ids,
+            txt=source_txt,
+            txt_ids=source_txt_ids,
+            y=source_vec,
+            timesteps=t_vec,
+            guidance=source_guidance_vec,
+            info=info
+        )
+        
+        info['inverse'] = False
+        v_tar = model(
+            img=z_tar,
+            img_ids=img_ids,
+            txt=target_txt,
+            txt_ids=target_txt_ids,
+            y=target_vec,
+            timesteps=t_vec,
+            guidance=target_guidance_vec,
+            info=info
+        )
+    
+        v_fe = v_tar - v_src[:, mask_indices,...]
+        z_fe = z_fe + (t_prev - t_curr) * v_fe * info['mask'][:, mask_indices,...]
+    return z_fe, info

flux/util.py

@@ -0,0 +1,201 @@
+import os
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from imwatermark import WatermarkEncoder
+from safetensors.torch import load_file as load_sft
+
+from flux.model import Flux, FluxParams
+from flux.modules.autoencoder import AutoEncoder, AutoEncoderParams
+from flux.modules.conditioner import HFEmbedder
+
+
+@dataclass
+class ModelSpec:
+    params: FluxParams
+    ae_params: AutoEncoderParams
+    ckpt_path: str | None
+    ae_path: str | None
+    repo_id: str | None
+    repo_flow: str | None
+    repo_ae: str | None
+
+configs = {
+    "flux-dev": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-dev",
+        repo_flow="flux1-dev.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_DEV"),
+        params=FluxParams(
+            in_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+            guidance_embed=True,
+        ),
+        ae_path=os.getenv("AE"),
+        ae_params=AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        ),
+    ),
+    "flux-schnell": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-schnell",
+        repo_flow="flux1-schnell.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_SCHNELL"),
+        params=FluxParams(
+            in_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+            guidance_embed=False,
+        ),
+        ae_path=os.getenv("AE"),
+        ae_params=AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        ),
+    ),
+}
+
+
+def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
+    if len(missing) > 0 and len(unexpected) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+        print("\n" + "-" * 79 + "\n")
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+    elif len(missing) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+    elif len(unexpected) > 0:
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+
+
+def load_flow_model(name: str, device: str | torch.device = "cuda", hf_download: bool = True, flux_cls=Flux) -> Flux:
+    # Loading Flux
+    print("Init model")
+    
+    ckpt_path = configs[name].ckpt_path
+    if (
+        ckpt_path is None
+        and configs[name].repo_id is not None
+        and configs[name].repo_flow is not None
+        and hf_download
+    ):
+        ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
+
+    with torch.device("meta" if ckpt_path is not None else device):
+        model = flux_cls(configs[name].params).to(torch.bfloat16)
+
+    if ckpt_path is not None:
+        print("Loading checkpoint")
+        # load_sft doesn't support torch.device
+        sd = load_sft(ckpt_path, device=str(device))
+        missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
+        print_load_warning(missing, unexpected)
+    return model
+
+
+def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder:
+    # max length 64, 128, 256 and 512 should work (if your sequence is short enough)
+    return HFEmbedder("google/t5-v1_1-xxl", max_length=max_length, is_clip=False, torch_dtype=torch.bfloat16).to(device)
+
+
+def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
+    return HFEmbedder("openai/clip-vit-large-patch14", max_length=77, is_clip=True, torch_dtype=torch.bfloat16).to(device)
+
+
+def load_ae(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder:
+    ckpt_path = configs[name].ae_path
+    if (
+        ckpt_path is None
+        and configs[name].repo_id is not None
+        and configs[name].repo_ae is not None
+        and hf_download
+    ):
+        ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_ae)
+
+    # Loading the autoencoder
+    print("Init AE")
+    with torch.device("meta" if ckpt_path is not None else device):
+        ae = AutoEncoder(configs[name].ae_params)
+
+    if ckpt_path is not None:
+        sd = load_sft(ckpt_path, device=str(device))
+        missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
+        print_load_warning(missing, unexpected)
+    return ae
+
+
+class WatermarkEmbedder:
+    def __init__(self, watermark):
+        self.watermark = watermark
+        self.num_bits = len(WATERMARK_BITS)
+        self.encoder = WatermarkEncoder()
+        self.encoder.set_watermark("bits", self.watermark)
+
+    def __call__(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Adds a predefined watermark to the input image
+
+        Args:
+            image: ([N,] B, RGB, H, W) in range [-1, 1]
+
+        Returns:
+            same as input but watermarked
+        """
+        image = 0.5 * image + 0.5
+        squeeze = len(image.shape) == 4
+        if squeeze:
+            image = image[None, ...]
+        n = image.shape[0]
+        image_np = rearrange((255 * image).detach().cpu(), "n b c h w -> (n b) h w c").numpy()[:, :, :, ::-1]
+        # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255]
+        # watermarking libary expects input as cv2 BGR format
+        for k in range(image_np.shape[0]):
+            image_np[k] = self.encoder.encode(image_np[k], "dwtDct")
+        image = torch.from_numpy(rearrange(image_np[:, :, :, ::-1], "(n b) h w c -> n b c h w", n=n)).to(
+            image.device
+        )
+        image = torch.clamp(image / 255, min=0.0, max=1.0)
+        if squeeze:
+            image = image[0]
+        image = 2 * image - 1
+        return image
+
+
+# A fixed 48-bit message that was chosen at random
+WATERMARK_MESSAGE = 0b001010101111111010000111100111001111010100101110
+# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
+WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]
+embed_watermark = WatermarkEmbedder(WATERMARK_BITS)

models/kv_edit.py

@@ -0,0 +1,213 @@
+from dataclasses import dataclass
+
+from einops import rearrange,repeat
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from typing import List
+
+from flux.sampling import get_schedule, unpack,denoise_kv,denoise_kv_inf
+from flux.util import load_flow_model
+from flux.model import Flux_kv
+
+@dataclass
+class SamplingOptions:
+    source_prompt: str = ''
+    target_prompt: str = ''
+    # prompt: str
+    width: int = 1366
+    height: int = 768
+    inversion_num_steps: int = 0
+    denoise_num_steps: int = 0
+    skip_step: int = 0
+    inversion_guidance: float = 1.0
+    denoise_guidance: float = 1.0
+    seed: int = 42
+    re_init: bool = False
+    attn_mask: bool = False
+
+class only_Flux(torch.nn.Module): # 仅包括初始化函数
+    def __init__(self, device,name='flux-dev'):
+        self.device = device
+        self.name = name
+        super().__init__()
+        self.model = load_flow_model(self.name, device=self.device,flux_cls=Flux_kv)
+        
+    def create_attention_mask(self,seq_len, mask_indices, text_len=512, device='cuda'):
+        """
+        创建自定义的注意力掩码。
+
+        Args:
+            seq_len (int): 序列长度。
+            mask_indices (List[int]): 图像令牌中掩码区域的索引。
+            text_len (int): 文本令牌的长度，默认 512。
+            device (str): 设备类型，如 'cuda' 或 'cpu'。
+
+        Returns:
+            torch.Tensor: 形状为 (seq_len, seq_len) 的注意力掩码。
+        """
+        # 初始化掩码为全 False
+        attention_mask = torch.zeros(seq_len, seq_len, dtype=torch.bool, device=device)
+
+        # 文本令牌索引
+        text_indices = torch.arange(0, text_len, device=device)
+
+        # 掩码区域令牌索引
+        mask_token_indices = torch.tensor([idx + text_len for idx in mask_indices], device=device)
+
+        # 背景区域令牌索引
+        all_indices = torch.arange(text_len, seq_len, device=device)
+        background_token_indices = torch.tensor([idx for idx in all_indices if idx not in mask_token_indices])
+
+        # 设置文本查询可以关注所有键
+        attention_mask[text_indices.unsqueeze(1).expand(-1, seq_len)] = True
+        attention_mask[text_indices.unsqueeze(1), text_indices] = True# 关注文本
+        attention_mask[text_indices.unsqueeze(1), background_token_indices] = True # 关注背景
+
+        
+        # attention_mask[mask_token_indices.unsqueeze(1), background_token_indices] = True  # 关注背景
+        attention_mask[mask_token_indices.unsqueeze(1), text_indices] = True  # 关注文本
+        attention_mask[mask_token_indices.unsqueeze(1), mask_token_indices] = True  # 关注掩码区域
+
+        
+        # attention_mask[background_token_indices.unsqueeze(1).expand(-1, seq_len), :] = False
+        attention_mask[background_token_indices.unsqueeze(1), mask_token_indices] = True  # 关注掩码
+        attention_mask[background_token_indices.unsqueeze(1), text_indices] = True  # 关注文本
+        attention_mask[background_token_indices.unsqueeze(1), background_token_indices] = True  # 关注背景区域
+
+        return attention_mask.unsqueeze(0)
+     
+class Flux_kv_edit_inf(only_Flux):
+    def __init__(self, device,name):
+        super().__init__(device,name)
+
+    @torch.inference_mode()
+    def forward(self,inp,inp_target,mask:Tensor,opts):
+        #############根据mask生成token序列上的索引试试#######################
+        info = {}
+        info['feature'] = {}
+        bs, L, d = inp["img"].shape
+        h = opts.height // 8
+        w = opts.width // 8
+        mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=False)
+        mask[mask > 0] = 1
+        
+        mask = repeat(mask, 'b c h w -> b (repeat c) h w', repeat=16)
+        # mask = F.max_pool2d(mask, kernel_size=3, stride=1, padding=1)
+        # mask = mask.flatten().to(self.device[1])
+        mask = rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+        info['mask'] = mask
+        bool_mask = (mask.sum(dim=2) > 0.5)
+        info['mask_indices'] = torch.nonzero(bool_mask)[:,1] # 使用花式索引 即 数字tensor索引tensor 这个是基于图像的 在seq中需要加512
+         #单独分离inversion
+        if opts.attn_mask and (~bool_mask).any(): # mask有一个false就进行attn mask 全true就none 
+            attention_mask = self.create_attention_mask(L+512, info['mask_indices'], device=self.device)
+        else:
+            attention_mask = None   
+        info['attention_mask'] = attention_mask
+        
+        denoise_timesteps = get_schedule(opts.denoise_num_steps, inp["img"].shape[1], shift=(self.name != "flux-schnell"))
+        # denoise_timesteps = get_schedule(opts.denoise_num_steps, inp_target["img"].shape[1], shift=False)
+        denoise_timesteps = denoise_timesteps[opts.skip_step:]
+    
+        z0 = inp["img"]
+
+        with torch.no_grad():
+            info['inject'] = True
+            z_fe, info = denoise_kv_inf(self.model, img=inp["img"], img_ids=inp['img_ids'], 
+                                    source_txt=inp['txt'], source_txt_ids=inp['txt_ids'], source_vec=inp['vec'],
+                                    target_txt=inp_target['txt'], target_txt_ids=inp_target['txt_ids'], target_vec=inp_target['vec'],
+                                    timesteps=denoise_timesteps, source_guidance=opts.inversion_guidance, target_guidance=opts.denoise_guidance,
+                                    info=info)
+        mask_indices = info['mask_indices'] # 图片seq坐标下的
+        # x是根据索引取出来的 再放回去
+        z0[:, mask_indices,...] = z_fe
+
+        # decode latents to pixel space
+        z0 = unpack(z0.float(),  opts.height, opts.width)
+        del info
+        return z0
+
+class Flux_kv_edit(only_Flux):
+    def __init__(self, device,name):
+        super().__init__(device,name)
+    
+    @torch.inference_mode()
+    def forward(self,inp,inp_target,mask:Tensor,opts):
+        z0,zt,info = self.inverse(inp,mask,opts)
+        z0 = self.denoise(z0,zt,inp_target,mask,opts,info)
+        return z0
+    @torch.inference_mode()
+    def inverse(self,inp,mask,opts):
+        info = {}
+        info['feature'] = {}
+        bs, L, d = inp["img"].shape
+        h = opts.height // 8
+        w = opts.width // 8
+        # mask = F.interpolate(mask, size=(h,w), mode='nearest')
+
+        if opts.attn_mask:
+            mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=False)
+            mask[mask > 0] = 1
+            
+            mask = repeat(mask, 'b c h w -> b (repeat c) h w', repeat=16)
+            # mask = F.max_pool2d(mask, kernel_size=3, stride=1, padding=1)
+            # mask = mask.flatten().to(self.device[1])
+            mask = rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+            bool_mask = (mask.sum(dim=2) > 0.5)
+            mask_indices = torch.nonzero(bool_mask)[:,1] # 使用花式索引 即 数字tensor索引tensor 这个是基于图像的 在seq中需要加512
+            
+            #单独分离inversion
+            assert not (~bool_mask).all(), "mask is all false"
+            assert not (bool_mask).all(), "mask is all true"
+            attention_mask = self.create_attention_mask(L+512, mask_indices, device=mask.device)
+            info['attention_mask'] = attention_mask
+    
+        
+        denoise_timesteps = get_schedule(opts.denoise_num_steps, inp["img"].shape[1], shift=(self.name != "flux-schnell"))
+        denoise_timesteps = denoise_timesteps[opts.skip_step:]
+        
+        # 加噪过程
+        z0 = inp["img"].clone()        
+        info['inverse'] = True
+        zt, info = denoise_kv(self.model, **inp, timesteps=denoise_timesteps, guidance=opts.inversion_guidance, inverse=True, info=info)
+        return z0,zt,info
+    
+    @torch.inference_mode()
+    def denoise(self,z0,zt,inp_target,mask:Tensor,opts,info):
+        
+        h = opts.height // 8
+        w = opts.width // 8
+        
+        mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=False)
+        mask[mask > 0] = 1
+        
+        mask = repeat(mask, 'b c h w -> b (repeat c) h w', repeat=16)
+      
+        mask = rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+        info['mask'] = mask
+        bool_mask = (mask.sum(dim=2) > 0.5)
+        info['mask_indices'] = torch.nonzero(bool_mask)[:,1] # 使用花式索引 即 数字tensor索引tensor 这个是基于图像的 在seq中需要加512
+        
+        denoise_timesteps = get_schedule(opts.denoise_num_steps, inp_target["img"].shape[1], shift=(self.name != "flux-schnell"))
+        denoise_timesteps = denoise_timesteps[opts.skip_step:]
+        # 重建的时候不需要全部token z这里需要根据indice拿出来
+        mask_indices = info['mask_indices'] # 图片seq坐标下的
+        if opts.re_init:
+            noise = torch.randn_like(zt)
+            t  = denoise_timesteps[0]
+            zt_noise = z0 *(1 - t) + noise * t
+            inp_target["img"] = zt_noise[:, mask_indices,...]
+        else:
+            inp_target["img"] = zt[:, mask_indices,...]
+
+        info['inverse'] = False
+        x, _ = denoise_kv(self.model, **inp_target, timesteps=denoise_timesteps, guidance=opts.denoise_guidance, inverse=False, info=info)
+        # x是根据索引取出来的 再放回去
+        z0[:, mask_indices,...] = z0[:, mask_indices,...] * (1 - info['mask'][:, mask_indices,...]) + x * info['mask'][:, mask_indices,...]
+        # x = inp['img'].clone()
+
+        # decode latents to pixel space
+        z0 = unpack(z0.float(),  opts.height, opts.width)
+        del info
+        return z0

requirements.txt

@@ -0,0 +1,18 @@
+torch
+einops
+accelerate==0.34.2
+einops==0.8.0
+transformers==4.41.2
+huggingface-hub==0.24.6
+datasets
+omegaconf
+diffusers
+sentencepiece
+opencv-python
+matplotlib
+onnxruntime
+torchvision
+timm
+invisible-watermark
+fire
+tqdm