Upload 5 files

uno/flux/math.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
+# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from einops import rearrange
+from torch import Tensor
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
+    q, k = apply_rope(q, k, pe)
+
+    x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+    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
+    omega = 1.0 / (theta**scale)
+    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)
+    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)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)

uno/flux/model.py

@@ -0,0 +1,222 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
+# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+
+from .modules.layers import DoubleStreamBlock, EmbedND, LastLayer, MLPEmbedder, SingleStreamBlock, 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.
+    """
+    _supports_gradient_checkpointing = True
+
+    def __init__(self, params: FluxParams):
+        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(
+            [
+                DoubleStreamBlock(
+                    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(
+            [
+                SingleStreamBlock(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)
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    @property
+    def attn_processors(self):
+        # set recursively
+        processors = {}  # type: dict[str, nn.Module]
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors):
+            if hasattr(module, "set_processor"):
+                processors[f"{name}.processor"] = module.processor
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    def set_attn_processor(self, processor):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    def forward(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+        y: Tensor,
+        guidance: Tensor | None = None,
+        ref_img: Tensor | None = None, 
+        ref_img_ids: 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)
+
+        # concat ref_img/img
+        img_end = img.shape[1]
+        if ref_img is not None:
+            if isinstance(ref_img, tuple) or isinstance(ref_img, list):
+                img_in = [img] + [self.img_in(ref) for ref in ref_img]
+                img_ids = [ids] + [ref_ids for ref_ids in ref_img_ids]
+                img = torch.cat(img_in, dim=1)  
+                ids = torch.cat(img_ids, dim=1)
+            else:
+                img = torch.cat((img, self.img_in(ref_img)), dim=1)  
+                ids = torch.cat((ids, ref_img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+        
+        for index_block, block in enumerate(self.double_blocks):
+            if self.training and self.gradient_checkpointing:
+                img, txt = torch.utils.checkpoint.checkpoint(
+                    block,
+                    img=img, 
+                    txt=txt, 
+                    vec=vec, 
+                    pe=pe, 
+                    use_reentrant=False,
+                )
+            else:
+                img, txt = block(
+                    img=img, 
+                    txt=txt, 
+                    vec=vec, 
+                    pe=pe
+                )
+
+        img = torch.cat((txt, img), 1)
+        for block in self.single_blocks:
+            if self.training and self.gradient_checkpointing:
+                img = torch.utils.checkpoint.checkpoint(
+                    block,
+                    img, vec=vec, pe=pe,
+                    use_reentrant=False
+                )
+            else:
+                img = block(img, vec=vec, pe=pe)
+        img = img[:, txt.shape[1] :, ...]
+        # index img
+        img = img[:, :img_end, ...]
+
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+        return img

uno/flux/pipeline.py

@@ -0,0 +1,322 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
+# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Literal
+
+import torch
+from einops import rearrange
+from PIL import ExifTags, Image
+import torchvision.transforms.functional as TVF
+
+from uno.flux.modules.layers import (
+    DoubleStreamBlockLoraProcessor,
+    DoubleStreamBlockProcessor,
+    SingleStreamBlockLoraProcessor,
+    SingleStreamBlockProcessor,
+)
+from uno.flux.sampling import denoise, get_noise, get_schedule, prepare_multi_ip, unpack
+from uno.flux.util import (
+    get_lora_rank,
+    load_ae,
+    load_checkpoint,
+    load_clip,
+    load_flow_model,
+    load_flow_model_only_lora,
+    load_flow_model_quintized,
+    load_t5,
+)
+
+
+def find_nearest_scale(image_h, image_w, predefined_scales):
+    """
+    根据图片的高度和宽度，找到最近的预定义尺度。
+
+    :param image_h: 图片的高度
+    :param image_w: 图片的宽度
+    :param predefined_scales: 预定义尺度列表 [(h1, w1), (h2, w2), ...]
+    :return: 最近的预定义尺度 (h, w)
+    """
+    # 计算输入图片的长宽比
+    image_ratio = image_h / image_w
+
+    # 初始化变量以存储最小差异和最近的尺度
+    min_diff = float('inf')
+    nearest_scale = None
+
+    # 遍历所有预定义尺度，找到与输入图片长宽比最接近的尺度
+    for scale_h, scale_w in predefined_scales:
+        predefined_ratio = scale_h / scale_w
+        diff = abs(predefined_ratio - image_ratio)
+
+        if diff < min_diff:
+            min_diff = diff
+            nearest_scale = (scale_h, scale_w)
+
+    return nearest_scale
+
+def preprocess_ref(raw_image: Image.Image, long_size: int = 512):
+    # 获取原始图像的宽度和高度
+    image_w, image_h = raw_image.size
+
+    # 计算长边和短边
+    if image_w >= image_h:
+        new_w = long_size
+        new_h = int((long_size / image_w) * image_h)
+    else:
+        new_h = long_size
+        new_w = int((long_size / image_h) * image_w)
+
+    # 按新的宽高进行等比例缩放
+    raw_image = raw_image.resize((new_w, new_h), resample=Image.LANCZOS)
+    target_w = new_w // 16 * 16
+    target_h = new_h // 16 * 16
+
+    # 计算裁剪的起始坐标以实现中心裁剪
+    left = (new_w - target_w) // 2
+    top = (new_h - target_h) // 2
+    right = left + target_w
+    bottom = top + target_h
+
+    # 进行中心裁剪
+    raw_image = raw_image.crop((left, top, right, bottom))
+
+    # 转换为 RGB 模式
+    raw_image = raw_image.convert("RGB")
+    return raw_image
+
+class UNOPipeline:
+    def __init__(
+        self,
+        model_type: str,
+        device: torch.device,
+        offload: bool = False,
+        only_lora: bool = False,
+        lora_rank: int = 16
+    ):
+        self.device = device
+        self.offload = offload
+        self.model_type = model_type
+
+        self.clip = load_clip(self.device)
+        self.t5 = load_t5(self.device, max_length=512)
+        self.ae = load_ae(model_type, device="cpu" if offload else self.device)
+        self.use_fp8 = "fp8" in model_type
+        if only_lora:
+            self.model = load_flow_model_only_lora(
+                model_type,
+                device="cpu" if offload else self.device,
+                lora_rank=lora_rank,
+                use_fp8=self.use_fp8
+            )
+        else:
+            self.model = load_flow_model(model_type, device="cpu" if offload else self.device)
+
+
+    def load_ckpt(self, ckpt_path):
+        if ckpt_path is not None:
+            from safetensors.torch import load_file as load_sft
+            print("Loading checkpoint to replace old keys")
+            # load_sft doesn't support torch.device
+            if ckpt_path.endswith('safetensors'):
+                sd = load_sft(ckpt_path, device='cpu')
+                missing, unexpected = self.model.load_state_dict(sd, strict=False, assign=True)
+            else:
+                dit_state = torch.load(ckpt_path, map_location='cpu')
+                sd = {}
+                for k in dit_state.keys():
+                    sd[k.replace('module.','')] = dit_state[k]
+                missing, unexpected = self.model.load_state_dict(sd, strict=False, assign=True)
+                self.model.to(str(self.device))
+            print(f"missing keys: {missing}\n\n\n\n\nunexpected keys: {unexpected}")
+
+    def set_lora(self, local_path: str = None, repo_id: str = None,
+                 name: str = None, lora_weight: int = 0.7):
+        checkpoint = load_checkpoint(local_path, repo_id, name)
+        self.update_model_with_lora(checkpoint, lora_weight)
+
+    def set_lora_from_collection(self, lora_type: str = "realism", lora_weight: int = 0.7):
+        checkpoint = load_checkpoint(
+            None, self.hf_lora_collection, self.lora_types_to_names[lora_type]
+        )
+        self.update_model_with_lora(checkpoint, lora_weight)
+
+    def update_model_with_lora(self, checkpoint, lora_weight):
+        rank = get_lora_rank(checkpoint)
+        lora_attn_procs = {}
+
+        for name, _ in self.model.attn_processors.items():
+            lora_state_dict = {}
+            for k in checkpoint.keys():
+                if name in k:
+                    lora_state_dict[k[len(name) + 1:]] = checkpoint[k] * lora_weight
+
+            if len(lora_state_dict):
+                if name.startswith("single_blocks"):
+                    lora_attn_procs[name] = SingleStreamBlockLoraProcessor(dim=3072, rank=rank)
+                else:
+                    lora_attn_procs[name] = DoubleStreamBlockLoraProcessor(dim=3072, rank=rank)
+                lora_attn_procs[name].load_state_dict(lora_state_dict)
+                lora_attn_procs[name].to(self.device)
+            else:
+                if name.startswith("single_blocks"):
+                    lora_attn_procs[name] = SingleStreamBlockProcessor()
+                else:
+                    lora_attn_procs[name] = DoubleStreamBlockProcessor()
+
+        self.model.set_attn_processor(lora_attn_procs)
+
+
+    def __call__(
+        self,
+        prompt: str,
+        width: int = 512,
+        height: int = 512,
+        guidance: float = 4,
+        num_steps: int = 50,
+        seed: int = 123456789,
+        **kwargs
+    ):
+        width = 16 * (width // 16)
+        height = 16 * (height // 16)
+
+        device_type = self.device if isinstance(self.device, str) else self.device.type
+        with torch.autocast(enabled=self.use_fp8, device_type=device_type, dtype=torch.bfloat16):
+            return self.forward(
+                prompt,
+                width,
+                height,
+                guidance,
+                num_steps,
+                seed,
+                **kwargs
+            )
+
+    @torch.inference_mode()
+    def gradio_generate(
+        self,
+        prompt: str,
+        width: int,
+        height: int,
+        guidance: float,
+        num_steps: int,
+        seed: int,
+        image_prompt1: Image.Image,
+        image_prompt2: Image.Image,
+        image_prompt3: Image.Image,
+        image_prompt4: Image.Image,
+    ):
+        ref_imgs = [image_prompt1, image_prompt2, image_prompt3, image_prompt4]
+        ref_imgs = [img for img in ref_imgs if isinstance(img, Image.Image)]
+        ref_long_side = 512 if len(ref_imgs) <= 1 else 320
+        ref_imgs = [preprocess_ref(img, ref_long_side) for img in ref_imgs]
+
+            # ✅ If seed is -1 (user wants random), sample a long int
+        if seed == -1 or seed is None:
+            seed = int(torch.randint(0, 2**63 - 1, (1,)).item())
+        else:
+            seed = int(seed)  # make sure it's a Python int
+
+        print(f"🧪 [DEBUG] Using seed: {seed} for image generation")
+
+        img = self(prompt=prompt, width=width, height=height, guidance=guidance,
+                   num_steps=num_steps, seed=seed, ref_imgs=ref_imgs)
+
+        filename = f"output/gradio/{seed}_{prompt[:20]}.png"
+        os.makedirs(os.path.dirname(filename), exist_ok=True)
+        exif_data = Image.Exif()
+        exif_data[ExifTags.Base.Make] = "UNO"
+        exif_data[ExifTags.Base.Model] = self.model_type
+        info = f"{prompt=}, {seed=}, {width=}, {height=}, {guidance=}, {num_steps=}"
+        exif_data[ExifTags.Base.ImageDescription] = info
+        img.save(filename, format="png", exif=exif_data)
+        return img, filename
+
+    @torch.inference_mode
+    def forward(
+        self,
+        prompt: str,
+        width: int,
+        height: int,
+        guidance: float,
+        num_steps: int,
+        seed: int,
+        ref_imgs: list[Image.Image] | None = None,
+        pe: Literal['d', 'h', 'w', 'o'] = 'd',
+    ):
+
+        # ✅ Ensure seed is always a valid integer
+        if seed == -1 or seed is None:
+            seed = int(torch.randint(0, 2**63 - 1, (1,)).item())
+        else:
+            seed = int(seed)
+        print(f"🧪 [DEBUG] Using seed: {seed}")
+
+        x = get_noise(
+            1, height, width, device=self.device,
+            dtype=torch.bfloat16, seed=seed
+        )
+        timesteps = get_schedule(
+            num_steps,
+            (width // 8) * (height // 8) // (16 * 16),
+            shift=True,
+        )
+        if self.offload:
+            self.ae.encoder = self.ae.encoder.to(self.device)
+        x_1_refs = [
+            self.ae.encode(
+                (TVF.to_tensor(ref_img) * 2.0 - 1.0) 
+                .unsqueeze(0).to(self.device, torch.float32)
+            ).to(torch.bfloat16)
+            for ref_img in ref_imgs
+        ]
+
+        if self.offload:
+            self.offload_model_to_cpu(self.ae.encoder)
+            self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
+        inp_cond = prepare_multi_ip(
+            t5=self.t5, clip=self.clip,
+            img=x,
+            prompt=prompt, ref_imgs=x_1_refs, pe=pe
+        )
+
+        if self.offload:
+            self.offload_model_to_cpu(self.t5, self.clip)
+            self.model = self.model.to(self.device)
+
+        x = denoise(
+            self.model,
+            **inp_cond,
+            timesteps=timesteps,
+            guidance=guidance,
+        )
+
+        if self.offload:
+            self.offload_model_to_cpu(self.model)
+            self.ae.decoder.to(x.device)
+        x = unpack(x.float(), height, width)
+        x = self.ae.decode(x)
+        self.offload_model_to_cpu(self.ae.decoder)
+
+        x1 = x.clamp(-1, 1)
+        x1 = rearrange(x1[-1], "c h w -> h w c")
+        output_img = Image.fromarray((127.5 * (x1 + 1.0)).cpu().byte().numpy())
+        return output_img
+
+    def offload_model_to_cpu(self, *models):
+        if not self.offload: return
+        for model in models:
+            model.cpu()
+            torch.cuda.empty_cache()

uno/flux/sampling.py

@@ -0,0 +1,252 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
+# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import Literal
+
+import torch
+from einops import rearrange, repeat
+from torch import Tensor
+from tqdm import tqdm
+
+from .model import Flux
+from .modules.conditioner import HFEmbedder
+
+
+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],
+    ref_img: None | Tensor=None,
+    pe: Literal['d', 'h', 'w', 'o'] ='d'
+) -> dict[str, Tensor]:
+    assert pe in ['d', 'h', 'w', 'o']
+    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 ref_img is not None:
+        _, _, ref_h, ref_w = ref_img.shape
+        ref_img = rearrange(ref_img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+        if ref_img.shape[0] == 1 and bs > 1:
+            ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)
+        ref_img_ids = torch.zeros(ref_h // 2, ref_w // 2, 3)
+        # img id分别在宽高偏移各自最大值
+        h_offset = h // 2 if pe in {'d', 'h'} else 0
+        w_offset = w // 2 if pe in {'d', 'w'} else 0
+        ref_img_ids[..., 1] = ref_img_ids[..., 1] + torch.arange(ref_h // 2)[:, None] + h_offset
+        ref_img_ids[..., 2] = ref_img_ids[..., 2] + torch.arange(ref_w // 2)[None, :] + w_offset
+        ref_img_ids = repeat(ref_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 ref_img is not None:
+        return {
+            "img": img,
+            "img_ids": img_ids.to(img.device),
+            "ref_img": ref_img,
+            "ref_img_ids": ref_img_ids.to(img.device),
+            "txt": txt.to(img.device),
+            "txt_ids": txt_ids.to(img.device),
+            "vec": vec.to(img.device),
+        }
+    else:
+        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_multi_ip(
+    t5: HFEmbedder,
+    clip: HFEmbedder,
+    img: Tensor,
+    prompt: str | list[str],
+    ref_imgs: list[Tensor] | None = None,
+    pe: Literal['d', 'h', 'w', 'o'] = 'd'
+) -> dict[str, Tensor]:
+    assert pe in ['d', 'h', 'w', 'o']
+    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)
+
+    ref_img_ids = []
+    ref_imgs_list = []
+    pe_shift_w, pe_shift_h = w // 2, h // 2
+    for ref_img in ref_imgs:
+        _, _, ref_h1, ref_w1 = ref_img.shape
+        ref_img = rearrange(ref_img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+        if ref_img.shape[0] == 1 and bs > 1:
+            ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)
+        ref_img_ids1 = torch.zeros(ref_h1 // 2, ref_w1 // 2, 3)
+        # img id分别���宽高偏移各自最大值
+        h_offset = pe_shift_h if pe in {'d', 'h'} else 0
+        w_offset = pe_shift_w if pe in {'d', 'w'} else 0
+        ref_img_ids1[..., 1] = ref_img_ids1[..., 1] + torch.arange(ref_h1 // 2)[:, None] + h_offset
+        ref_img_ids1[..., 2] = ref_img_ids1[..., 2] + torch.arange(ref_w1 // 2)[None, :] + w_offset
+        ref_img_ids1 = repeat(ref_img_ids1, "h w c -> b (h w) c", b=bs)
+        ref_img_ids.append(ref_img_ids1)
+        ref_imgs_list.append(ref_img)
+
+        # 更新pe shift
+        pe_shift_h += ref_h1 // 2
+        pe_shift_w += ref_w1 // 2
+
+    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),
+        "ref_img": tuple(ref_imgs_list),
+        "ref_img_ids": [ref_img_id.to(img.device) for ref_img_id in ref_img_ids],
+        "txt": txt.to(img.device),
+        "txt_ids": txt_ids.to(img.device),
+        "vec": 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
+):
+    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:
+        # eastimate 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,
+    ref_img: Tensor=None,
+    ref_img_ids: Tensor=None,
+):
+    i = 0
+    guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
+    for t_curr, t_prev in tqdm(zip(timesteps[:-1], timesteps[1:]), total=len(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,
+            ref_img=ref_img,
+            ref_img_ids=ref_img_ids,
+            txt=txt,
+            txt_ids=txt_ids,
+            y=vec,
+            timesteps=t_vec,
+            guidance=guidance_vec
+        )
+        img = img + (t_prev - t_curr) * pred
+        i += 1
+    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,
+    )

uno/flux/util.py

@@ -0,0 +1,411 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
+# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from dataclasses import dataclass
+
+import torch
+import json
+import numpy as np
+from huggingface_hub import hf_hub_download
+from safetensors import safe_open
+from safetensors.torch import load_file as load_sft
+
+from .model import Flux, FluxParams
+from .modules.autoencoder import AutoEncoder, AutoEncoderParams
+from .modules.conditioner import HFEmbedder
+
+import re
+from uno.flux.modules.layers import DoubleStreamBlockLoraProcessor, SingleStreamBlockLoraProcessor
+def load_model(ckpt, device='cpu'):
+    if ckpt.endswith('safetensors'):
+        from safetensors import safe_open
+        pl_sd = {}
+        with safe_open(ckpt, framework="pt", device=device) as f:
+            for k in f.keys():
+                pl_sd[k] = f.get_tensor(k)
+    else:
+        pl_sd = torch.load(ckpt, map_location=device)
+    return pl_sd
+
+def load_safetensors(path):
+    tensors = {}
+    with safe_open(path, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            tensors[key] = f.get_tensor(key)
+    return tensors
+
+def get_lora_rank(checkpoint):
+    for k in checkpoint.keys():
+        if k.endswith(".down.weight"):
+            return checkpoint[k].shape[0]
+
+def load_checkpoint(local_path, repo_id, name):
+    if local_path is not None:
+        if '.safetensors' in local_path:
+            print(f"Loading .safetensors checkpoint from {local_path}")
+            checkpoint = load_safetensors(local_path)
+        else:
+            print(f"Loading checkpoint from {local_path}")
+            checkpoint = torch.load(local_path, map_location='cpu')
+    elif repo_id is not None and name is not None:
+        print(f"Loading checkpoint {name} from repo id {repo_id}")
+        checkpoint = load_from_repo_id(repo_id, name)
+    else:
+        raise ValueError(
+            "LOADING ERROR: you must specify local_path or repo_id with name in HF to download"
+        )
+    return checkpoint
+
+
+def c_crop(image):
+    width, height = image.size
+    new_size = min(width, height)
+    left = (width - new_size) / 2
+    top = (height - new_size) / 2
+    right = (width + new_size) / 2
+    bottom = (height + new_size) / 2
+    return image.crop((left, top, right, bottom))
+
+def pad64(x):
+    return int(np.ceil(float(x) / 64.0) * 64 - x)
+
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
+
+@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
+    repo_id_ae: str | None
+
+
+configs = {
+    "flux-dev": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-dev",
+        repo_id_ae="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-dev-fp8": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-dev",
+        repo_id_ae="black-forest-labs/FLUX.1-dev",
+        repo_flow="flux1-dev.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_DEV_FP8"),
+        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_id_ae="black-forest-labs/FLUX.1-dev",
+        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_from_repo_id(repo_id, checkpoint_name):
+    ckpt_path = hf_hub_download(repo_id, checkpoint_name)
+    sd = load_sft(ckpt_path, device='cpu')
+    return sd
+
+def load_flow_model(name: str, device: str | torch.device = "cuda", hf_download: bool = True):
+    # 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(configs[name].params).to(torch.bfloat16)
+
+    if ckpt_path is not None:
+        print("Loading checkpoint")
+        # load_sft doesn't support torch.device
+        sd = load_model(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_flow_model_only_lora(
+    name: str,
+    device: str | torch.device = "cuda",
+    hf_download: bool = True,
+    lora_rank: int = 16,
+    use_fp8: bool = False
+):
+    # 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.replace("sft", "safetensors"))
+    
+    if hf_download:
+        try:
+            lora_ckpt_path = hf_hub_download("bytedance-research/UNO", "dit_lora.safetensors")
+        except:
+            lora_ckpt_path = os.environ.get("LORA", None)
+    else:
+        lora_ckpt_path = os.environ.get("LORA", None)
+
+    with torch.device("meta" if ckpt_path is not None else device):
+        model = Flux(configs[name].params)
+
+
+    model = set_lora(model, lora_rank, device="meta" if lora_ckpt_path is not None else device)
+
+    if ckpt_path is not None:
+        print("Loading lora")
+        lora_sd = load_sft(lora_ckpt_path, device=str(device)) if lora_ckpt_path.endswith("safetensors")\
+            else torch.load(lora_ckpt_path, map_location='cpu')
+        
+        print("Loading main checkpoint")
+        # load_sft doesn't support torch.device
+
+        if ckpt_path.endswith('safetensors'):
+            if use_fp8:
+                print(
+                    "####\n"
+                    "We are in fp8 mode right now, since the fp8 checkpoint of XLabs-AI/flux-dev-fp8 seems broken\n"
+                    "we convert the fp8 checkpoint on flight from bf16 checkpoint\n"
+                    "If your storage is constrained"
+                    "you can save the fp8 checkpoint and replace the bf16 checkpoint by yourself\n"
+                )
+                sd = load_sft(ckpt_path, device="cpu")
+                sd = {k: v.to(dtype=torch.float8_e4m3fn, device=device) for k, v in sd.items()}
+            else:
+                sd = load_sft(ckpt_path, device=str(device))
+            
+            sd.update(lora_sd)
+            missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
+        else:
+            dit_state = torch.load(ckpt_path, map_location='cpu')
+            sd = {}
+            for k in dit_state.keys():
+                sd[k.replace('module.','')] = dit_state[k]
+            sd.update(lora_sd)
+            missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
+            model.to(str(device))
+        print_load_warning(missing, unexpected)
+    return model
+
+
+def set_lora(
+    model: Flux,
+    lora_rank: int,
+    double_blocks_indices: list[int] | None = None,
+    single_blocks_indices: list[int] | None = None,
+    device: str | torch.device = "cpu",
+) -> Flux:
+    double_blocks_indices = list(range(model.params.depth)) if double_blocks_indices is None else double_blocks_indices
+    single_blocks_indices = list(range(model.params.depth_single_blocks)) if single_blocks_indices is None \
+                            else single_blocks_indices
+    
+    lora_attn_procs = {}
+    with torch.device(device):
+        for name, attn_processor in  model.attn_processors.items():
+            match = re.search(r'\.(\d+)\.', name)
+            if match:
+                layer_index = int(match.group(1))
+
+            if name.startswith("double_blocks") and layer_index in double_blocks_indices:
+                lora_attn_procs[name] = DoubleStreamBlockLoraProcessor(dim=model.params.hidden_size, rank=lora_rank)
+            elif name.startswith("single_blocks") and layer_index in single_blocks_indices:
+                lora_attn_procs[name] = SingleStreamBlockLoraProcessor(dim=model.params.hidden_size, rank=lora_rank)
+            else:
+                lora_attn_procs[name] = attn_processor
+    model.set_attn_processor(lora_attn_procs)
+    return model
+
+
+def load_flow_model_quintized(name: str, device: str | torch.device = "cuda", hf_download: bool = True):
+    # Loading Flux
+    from optimum.quanto import requantize
+    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)
+    # json_path = hf_hub_download(configs[name].repo_id, 'flux_dev_quantization_map.json')
+
+
+    model = Flux(configs[name].params).to(torch.bfloat16)
+
+    print("Loading checkpoint")
+    # load_sft doesn't support torch.device
+    sd = load_sft(ckpt_path, device='cpu')
+    sd = {k: v.to(dtype=torch.float8_e4m3fn, device=device) for k, v in sd.items()}
+    model.load_state_dict(sd, assign=True)
+    return model
+    with open(json_path, "r") as f:
+        quantization_map = json.load(f)
+    print("Start a quantization process...")
+    requantize(model, sd, quantization_map, device=device)
+    print("Model is quantized!")
+    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)
+    version = os.environ.get("T5", "xlabs-ai/xflux_text_encoders")
+    return HFEmbedder(version, max_length=max_length, torch_dtype=torch.bfloat16).to(device)
+
+def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
+    version = os.environ.get("CLIP", "openai/clip-vit-large-patch14")
+    return HFEmbedder(version, max_length=77, 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_ae, 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