# -*- coding: utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. # This file contains code that is adapted from # https://github.com/black-forest-labs/flux.git import math import torch from torch import Tensor, nn from collections import OrderedDict from functools import partial from einops import rearrange, repeat from scepter.modules.model.base_model import BaseModel from scepter.modules.model.registry import BACKBONES from scepter.modules.utils.config import dict_to_yaml from scepter.modules.utils.distribute import we from scepter.modules.utils.file_system import FS from torch.utils.checkpoint import checkpoint_sequential from torch.nn.utils.rnn import pad_sequence from .layers import (DoubleStreamBlock, EmbedND, LastLayer, MLPEmbedder, SingleStreamBlock, timestep_embedding) @BACKBONES.register_class() class Flux(BaseModel): """ Transformer backbone Diffusion model with RoPE. """ para_dict = { 'IN_CHANNELS': { 'value': 64, 'description': "model's input channels." }, 'OUT_CHANNELS': { 'value': 64, 'description': "model's output channels." }, 'HIDDEN_SIZE': { 'value': 1024, 'description': "model's hidden size." }, 'NUM_HEADS': { 'value': 16, 'description': 'number of heads in the transformer.' }, 'AXES_DIM': { 'value': [16, 56, 56], 'description': 'dimensions of the axes of the positional encoding.' }, 'THETA': { 'value': 10_000, 'description': 'theta for positional encoding.' }, 'VEC_IN_DIM': { 'value': 768, 'description': 'dimension of the vector input.' }, 'GUIDANCE_EMBED': { 'value': False, 'description': 'whether to use guidance embedding.' }, 'CONTEXT_IN_DIM': { 'value': 4096, 'description': 'dimension of the context input.' }, 'MLP_RATIO': { 'value': 4.0, 'description': 'ratio of mlp hidden size to hidden size.' }, 'QKV_BIAS': { 'value': True, 'description': 'whether to use bias in qkv projection.' }, 'DEPTH': { 'value': 19, 'description': 'number of transformer blocks.' }, 'DEPTH_SINGLE_BLOCKS': { 'value': 38, 'description': 'number of transformer blocks in the single stream block.' }, 'USE_GRAD_CHECKPOINT': { 'value': False, 'description': 'whether to use gradient checkpointing.' } } def __init__(self, cfg, logger=None): super().__init__(cfg, logger=logger) self.in_channels = cfg.IN_CHANNELS self.out_channels = cfg.get('OUT_CHANNELS', self.in_channels) hidden_size = cfg.get('HIDDEN_SIZE', 1024) num_heads = cfg.get('NUM_HEADS', 16) axes_dim = cfg.AXES_DIM theta = cfg.THETA vec_in_dim = cfg.VEC_IN_DIM self.guidance_embed = cfg.GUIDANCE_EMBED context_in_dim = cfg.CONTEXT_IN_DIM mlp_ratio = cfg.MLP_RATIO qkv_bias = cfg.QKV_BIAS depth = cfg.DEPTH depth_single_blocks = cfg.DEPTH_SINGLE_BLOCKS self.use_grad_checkpoint = cfg.get("USE_GRAD_CHECKPOINT", False) self.attn_backend = cfg.get("ATTN_BACKEND", "pytorch") self.cache_pretrain_model = cfg.get("CACHE_PRETRAIN_MODEL", False) self.lora_model = cfg.get("DIFFUSERS_LORA_MODEL", None) self.comfyui_lora_model = cfg.get("COMFYUI_LORA_MODEL", None) self.swift_lora_model = cfg.get("SWIFT_LORA_MODEL", None) self.blackforest_lora_model = cfg.get("BLACKFOREST_LORA_MODEL", None) self.pretrain_adapter = cfg.get("PRETRAIN_ADAPTER", None) if hidden_size % num_heads != 0: raise ValueError( f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}" ) pe_dim = hidden_size // num_heads if sum(axes_dim) != pe_dim: raise ValueError( f"Got {axes_dim} but expected positional dim {pe_dim}") self.hidden_size = hidden_size self.num_heads = num_heads self.pe_embedder = EmbedND(dim=pe_dim, theta=theta, axes_dim=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(vec_in_dim, self.hidden_size) self.guidance_in = (MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if self.guidance_embed else nn.Identity()) self.txt_in = nn.Linear(context_in_dim, self.hidden_size) self.double_blocks = nn.ModuleList( [ DoubleStreamBlock( self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, backend=self.attn_backend ) for _ in range(depth) ] ) self.single_blocks = nn.ModuleList( [ SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, backend=self.attn_backend) for _ in range(depth_single_blocks) ] ) self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels) def prepare_input(self, x, context, y, x_shape=None): # x.shape [6, 16, 16, 16] target is [6, 16, 768, 1360] bs, c, h, w = x.shape x = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) x_id = torch.zeros(h // 2, w // 2, 3) x_id[..., 1] = x_id[..., 1] + torch.arange(h // 2)[:, None] x_id[..., 2] = x_id[..., 2] + torch.arange(w // 2)[None, :] x_ids = repeat(x_id, "h w c -> b (h w) c", b=bs) txt_ids = torch.zeros(bs, context.shape[1], 3) return x, x_ids.to(x), context.to(x), txt_ids.to(x), y.to(x), h, w def unpack(self, 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/2), w=math.ceil(width/2), ph=2, pw=2, ) def merge_diffuser_lora(self, ori_sd, lora_sd, scale=1.0): key_map = { "single_blocks.{}.linear1.weight": {"key_list": [ ["transformer.single_transformer_blocks.{}.attn.to_q.lora_A.weight", "transformer.single_transformer_blocks.{}.attn.to_q.lora_B.weight", [0, 3072]], ["transformer.single_transformer_blocks.{}.attn.to_k.lora_A.weight", "transformer.single_transformer_blocks.{}.attn.to_k.lora_B.weight", [3072, 6144]], ["transformer.single_transformer_blocks.{}.attn.to_v.lora_A.weight", "transformer.single_transformer_blocks.{}.attn.to_v.lora_B.weight", [6144, 9216]], ["transformer.single_transformer_blocks.{}.proj_mlp.lora_A.weight", "transformer.single_transformer_blocks.{}.proj_mlp.lora_B.weight", [9216, 21504]] ], "num": 38}, "single_blocks.{}.modulation.lin.weight": {"key_list": [ ["transformer.single_transformer_blocks.{}.norm.linear.lora_A.weight", "transformer.single_transformer_blocks.{}.norm.linear.lora_B.weight", [0, 9216]], ], "num": 38}, "single_blocks.{}.linear2.weight": {"key_list": [ ["transformer.single_transformer_blocks.{}.proj_out.lora_A.weight", "transformer.single_transformer_blocks.{}.proj_out.lora_B.weight", [0, 3072]], ], "num": 38}, "double_blocks.{}.txt_attn.qkv.weight": {"key_list": [ ["transformer.transformer_blocks.{}.attn.add_q_proj.lora_A.weight", "transformer.transformer_blocks.{}.attn.add_q_proj.lora_B.weight", [0, 3072]], ["transformer.transformer_blocks.{}.attn.add_k_proj.lora_A.weight", "transformer.transformer_blocks.{}.attn.add_k_proj.lora_B.weight", [3072, 6144]], ["transformer.transformer_blocks.{}.attn.add_v_proj.lora_A.weight", "transformer.transformer_blocks.{}.attn.add_v_proj.lora_B.weight", [6144, 9216]], ], "num": 19}, "double_blocks.{}.img_attn.qkv.weight": {"key_list": [ ["transformer.transformer_blocks.{}.attn.to_q.lora_A.weight", "transformer.transformer_blocks.{}.attn.to_q.lora_B.weight", [0, 3072]], ["transformer.transformer_blocks.{}.attn.to_k.lora_A.weight", "transformer.transformer_blocks.{}.attn.to_k.lora_B.weight", [3072, 6144]], ["transformer.transformer_blocks.{}.attn.to_v.lora_A.weight", "transformer.transformer_blocks.{}.attn.to_v.lora_B.weight", [6144, 9216]], ], "num": 19}, "double_blocks.{}.img_attn.proj.weight": {"key_list": [ ["transformer.transformer_blocks.{}.attn.to_out.0.lora_A.weight", "transformer.transformer_blocks.{}.attn.to_out.0.lora_B.weight", [0, 3072]] ], "num": 19}, "double_blocks.{}.txt_attn.proj.weight": {"key_list": [ ["transformer.transformer_blocks.{}.attn.to_add_out.lora_A.weight", "transformer.transformer_blocks.{}.attn.to_add_out.lora_B.weight", [0, 3072]] ], "num": 19}, "double_blocks.{}.img_mlp.0.weight": {"key_list": [ ["transformer.transformer_blocks.{}.ff.net.0.proj.lora_A.weight", "transformer.transformer_blocks.{}.ff.net.0.proj.lora_B.weight", [0, 12288]] ], "num": 19}, "double_blocks.{}.img_mlp.2.weight": {"key_list": [ ["transformer.transformer_blocks.{}.ff.net.2.lora_A.weight", "transformer.transformer_blocks.{}.ff.net.2.lora_B.weight", [0, 3072]] ], "num": 19}, "double_blocks.{}.txt_mlp.0.weight": {"key_list": [ ["transformer.transformer_blocks.{}.ff_context.net.0.proj.lora_A.weight", "transformer.transformer_blocks.{}.ff_context.net.0.proj.lora_B.weight", [0, 12288]] ], "num": 19}, "double_blocks.{}.txt_mlp.2.weight": {"key_list": [ ["transformer.transformer_blocks.{}.ff_context.net.2.lora_A.weight", "transformer.transformer_blocks.{}.ff_context.net.2.lora_B.weight", [0, 3072]] ], "num": 19}, "double_blocks.{}.img_mod.lin.weight": {"key_list": [ ["transformer.transformer_blocks.{}.norm1.linear.lora_A.weight", "transformer.transformer_blocks.{}.norm1.linear.lora_B.weight", [0, 18432]] ], "num": 19}, "double_blocks.{}.txt_mod.lin.weight": {"key_list": [ ["transformer.transformer_blocks.{}.norm1_context.linear.lora_A.weight", "transformer.transformer_blocks.{}.norm1_context.linear.lora_B.weight", [0, 18432]] ], "num": 19} } cover_lora_keys = set() cover_ori_keys = set() for k, v in key_map.items(): key_list = v["key_list"] block_num = v["num"] for block_id in range(block_num): for k_list in key_list: if k_list[0].format(block_id) in lora_sd and k_list[1].format(block_id) in lora_sd: cover_lora_keys.add(k_list[0].format(block_id)) cover_lora_keys.add(k_list[1].format(block_id)) current_weight = torch.matmul(lora_sd[k_list[0].format(block_id)].permute(1, 0), lora_sd[k_list[1].format(block_id)].permute(1, 0)).permute(1, 0) ori_sd[k.format(block_id)][k_list[2][0]:k_list[2][1], ...] += scale * current_weight cover_ori_keys.add(k.format(block_id)) # lora_sd.pop(k_list[0].format(block_id)) # lora_sd.pop(k_list[1].format(block_id)) self.logger.info(f"merge_blackforest_lora loads lora'parameters lora-paras: \n" f"cover-{len(cover_lora_keys)} vs total {len(lora_sd)} \n" f"cover ori-{len(cover_ori_keys)} vs total {len(ori_sd)}") return ori_sd def merge_swift_lora(self, ori_sd, lora_sd, scale = 1.0): have_lora_keys = {} for k, v in lora_sd.items(): k = k[len("model."):] if k.startswith("model.") else k ori_key = k.split("lora")[0] + "weight" if ori_key not in ori_sd: raise f"{ori_key} should in the original statedict" if ori_key not in have_lora_keys: have_lora_keys[ori_key] = {} if "lora_A" in k: have_lora_keys[ori_key]["lora_A"] = v elif "lora_B" in k: have_lora_keys[ori_key]["lora_B"] = v else: raise NotImplementedError self.logger.info(f"merge_swift_lora loads lora'parameters {len(have_lora_keys)}") for key, v in have_lora_keys.items(): current_weight = torch.matmul(v["lora_A"].permute(1, 0), v["lora_B"].permute(1, 0)).permute(1, 0) ori_sd[key] += scale * current_weight return ori_sd def merge_blackforest_lora(self, ori_sd, lora_sd, scale = 1.0): have_lora_keys = {} cover_lora_keys = set() cover_ori_keys = set() for k, v in lora_sd.items(): if "lora" in k: ori_key = k.split("lora")[0] + "weight" if ori_key not in ori_sd: raise f"{ori_key} should in the original statedict" if ori_key not in have_lora_keys: have_lora_keys[ori_key] = {} if "lora_A" in k: have_lora_keys[ori_key]["lora_A"] = v cover_lora_keys.add(k) cover_ori_keys.add(ori_key) elif "lora_B" in k: have_lora_keys[ori_key]["lora_B"] = v cover_lora_keys.add(k) cover_ori_keys.add(ori_key) else: if k in ori_sd: ori_sd[k] = v cover_lora_keys.add(k) cover_ori_keys.add(k) else: print("unsurpport keys: ", k) self.logger.info(f"merge_blackforest_lora loads lora'parameters lora-paras: \n" f"cover-{len(cover_lora_keys)} vs total {len(lora_sd)} \n" f"cover ori-{len(cover_ori_keys)} vs total {len(ori_sd)}") for key, v in have_lora_keys.items(): current_weight = torch.matmul(v["lora_A"].permute(1, 0), v["lora_B"].permute(1, 0)).permute(1, 0) # print(key, ori_sd[key].shape, current_weight.shape) ori_sd[key] += scale * current_weight return ori_sd def merge_comfyui_lora(self, ori_sd, lora_sd, scale = 1.0): ori_key_map = {key.replace("_", ".") : key for key in ori_sd.keys()} parse_ckpt = OrderedDict() for k, v in lora_sd.items(): if "alpha" in k: continue k = k.replace("lora_unet_", "").replace("_", ".") map_k = ori_key_map[k.split(".lora")[0] + ".weight"] if map_k not in parse_ckpt: parse_ckpt[map_k] = {} if "lora.up" in k: parse_ckpt[map_k]["lora_up"] = v elif "lora.down" in k: parse_ckpt[map_k]["lora_down"] = v if self.cache_pretrain_model: self.lora_dict[self.comfyui_lora_model] = {} for key, v in parse_ckpt.items(): current_weight = torch.matmul(v["lora_down"].permute(1, 0), v["lora_up"].permute(1, 0)).permute(1, 0) self.lora_dict[self.comfyui_lora_model] = current_weight ori_sd[key] += scale * current_weight return ori_sd def easy_lora_merge(self, ori_sd, lora_sd, scale = 1.0): for key, v in lora_sd.items(): ori_sd[key] += scale * v return ori_sd def load_pretrained_model(self, pretrained_model, lora_scale = 1.0): if next(self.parameters()).device.type == 'meta': map_location = torch.device(we.device_id) safe_device = we.device_id else: map_location = "cpu" safe_device = "cpu" if pretrained_model is not None: if not hasattr(self, "ckpt"): with FS.get_from(pretrained_model, wait_finish=True) as local_model: if local_model.endswith('safetensors'): from safetensors.torch import load_file as load_safetensors ckpt = load_safetensors(local_model, device=safe_device) else: ckpt = torch.load(local_model, map_location=map_location, weights_only=True) if "state_dict" in ckpt: ckpt = ckpt["state_dict"] if "model" in ckpt: ckpt = ckpt["model"]["model"] if self.cache_pretrain_model: self.ckpt = ckpt self.lora_dict = {} else: ckpt = self.ckpt new_ckpt = OrderedDict() for k, v in ckpt.items(): if k in ("img_in.weight"): model_p = self.state_dict()[k] if v.shape != model_p.shape: expanded_state_dict_weight = torch.zeros_like(model_p, device=v.device) slices = tuple(slice(0, dim) for dim in v.shape) expanded_state_dict_weight[slices] = v new_ckpt[k] = expanded_state_dict_weight else: new_ckpt[k] = v else: new_ckpt[k] = v if self.lora_model is not None: with FS.get_from(self.lora_model, wait_finish=True) as local_model: if local_model.endswith('safetensors'): from safetensors.torch import load_file as load_safetensors lora_sd = load_safetensors(local_model, device=safe_device) else: lora_sd = torch.load(local_model, map_location=map_location, weights_only=True) new_ckpt = self.merge_diffuser_lora(new_ckpt, lora_sd, scale=lora_scale) if self.swift_lora_model is not None: if not isinstance(self.swift_lora_model, list): self.swift_lora_model = [(self.swift_lora_model, 1.0)] for lora_model in self.swift_lora_model: if isinstance(lora_model, str): lora_model = (lora_model, 1.0/len(self.swift_lora_model)) print(lora_model) self.logger.info(f"load swift lora model: {lora_model}") with FS.get_from(lora_model[0], wait_finish=True) as local_model: if local_model.endswith('safetensors'): from safetensors.torch import load_file as load_safetensors lora_sd = load_safetensors(local_model, device=safe_device) else: lora_sd = torch.load(local_model, map_location=map_location, weights_only=True) new_ckpt = self.merge_swift_lora(new_ckpt, lora_sd, scale=lora_model[1]) if self.blackforest_lora_model is not None: with FS.get_from(self.blackforest_lora_model, wait_finish=True) as local_model: if local_model.endswith('safetensors'): from safetensors.torch import load_file as load_safetensors lora_sd = load_safetensors(local_model, device=safe_device) else: lora_sd = torch.load(local_model, map_location=map_location, weights_only=True) new_ckpt = self.merge_blackforest_lora(new_ckpt, lora_sd, scale=lora_scale) if self.comfyui_lora_model is not None: if hasattr(self, "current_lora") and self.current_lora == self.comfyui_lora_model: return if hasattr(self, "lora_dict") and self.comfyui_lora_model in self.lora_dict: new_ckpt = self.easy_lora_merge(new_ckpt, self.lora_dict[self.comfyui_lora_model], scale=lora_scale) else: with FS.get_from(self.comfyui_lora_model, wait_finish=True) as local_model: if local_model.endswith('safetensors'): from safetensors.torch import load_file as load_safetensors lora_sd = load_safetensors(local_model, device=safe_device) else: lora_sd = torch.load(local_model, map_location=map_location, weights_only=True) new_ckpt = self.merge_comfyui_lora(new_ckpt, lora_sd, scale=lora_scale) if self.comfyui_lora_model: self.current_lora = self.comfyui_lora_model adapter_ckpt = {} if self.pretrain_adapter is not None: with FS.get_from(self.pretrain_adapter, wait_finish=True) as local_adapter: if local_adapter.endswith('safetensors'): from safetensors.torch import load_file as load_safetensors adapter_ckpt = load_safetensors(local_adapter, device=safe_device) else: adapter_ckpt = torch.load(local_adapter, map_location=map_location, weights_only=True) new_ckpt.update(adapter_ckpt) missing, unexpected = self.load_state_dict(new_ckpt, strict=False, assign=True) self.logger.info( f'Restored from {pretrained_model} with {len(missing)} missing and {len(unexpected)} unexpected keys' ) if len(missing) > 0: self.logger.info(f'Missing Keys:\n {missing}') if len(unexpected) > 0: self.logger.info(f'\nUnexpected Keys:\n {unexpected}') def forward( self, x: Tensor, t: Tensor, cond: dict = {}, guidance: Tensor | None = None, gc_seg: int = 0 ) -> Tensor: x, x_ids, txt, txt_ids, y, h, w = self.prepare_input(x, cond["context"], cond["y"]) # running on sequences img x = self.img_in(x) vec = self.time_in(timestep_embedding(t, 256)) if self.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, x_ids), dim=1) pe = self.pe_embedder(ids) kwargs = dict( vec=vec, pe=pe, txt_length=txt.shape[1], ) x = torch.cat((txt, x), 1) if self.use_grad_checkpoint and gc_seg >= 0: x = checkpoint_sequential( functions=[partial(block, **kwargs) for block in self.double_blocks], segments=gc_seg if gc_seg > 0 else len(self.double_blocks), input=x, use_reentrant=False ) else: for block in self.double_blocks: x = block(x, **kwargs) kwargs = dict( vec=vec, pe=pe, ) if self.use_grad_checkpoint and gc_seg >= 0: x = checkpoint_sequential( functions=[partial(block, **kwargs) for block in self.single_blocks], segments=gc_seg if gc_seg > 0 else len(self.single_blocks), input=x, use_reentrant=False ) else: for block in self.single_blocks: x = block(x, **kwargs) x = x[:, txt.shape[1] :, ...] x = self.final_layer(x, vec) # (N, T, patch_size ** 2 * out_channels) 6 64 64 x = self.unpack(x, h, w) return x @staticmethod def get_config_template(): return dict_to_yaml('BACKBONE', __class__.__name__, Flux.para_dict, set_name=True) @BACKBONES.register_class() class FluxMR(Flux): def prepare_input(self, x, cond): if isinstance(cond['context'], list): context, y = torch.cat(cond["context"], dim=0).to(x), torch.cat(cond["y"], dim=0).to(x) else: context, y = cond['context'].to(x), cond['y'].to(x) batch_frames, batch_frames_ids = [], [] for ix, shape in zip(x, cond["x_shapes"]): # unpack image from sequence ix = ix[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1]) c, h, w = ix.shape ix = rearrange(ix, "c (h ph) (w pw) -> (h w) (c ph pw)", ph=2, pw=2) ix_id = torch.zeros(h // 2, w // 2, 3) ix_id[..., 1] = ix_id[..., 1] + torch.arange(h // 2)[:, None] ix_id[..., 2] = ix_id[..., 2] + torch.arange(w // 2)[None, :] ix_id = rearrange(ix_id, "h w c -> (h w) c") batch_frames.append([ix]) batch_frames_ids.append([ix_id]) x_list, x_id_list, mask_x_list, x_seq_length = [], [], [], [] for frames, frame_ids in zip(batch_frames, batch_frames_ids): proj_frames = [] for idx, one_frame in enumerate(frames): one_frame = self.img_in(one_frame) proj_frames.append(one_frame) ix = torch.cat(proj_frames, dim=0) if_id = torch.cat(frame_ids, dim=0) x_list.append(ix) x_id_list.append(if_id) mask_x_list.append(torch.ones(ix.shape[0]).to(ix.device, non_blocking=True).bool()) x_seq_length.append(ix.shape[0]) x = pad_sequence(tuple(x_list), batch_first=True) x_ids = pad_sequence(tuple(x_id_list), batch_first=True).to(x) # [b,pad_seq,2] pad (0.,0.) at dim2 mask_x = pad_sequence(tuple(mask_x_list), batch_first=True) txt = self.txt_in(context) txt_ids = torch.zeros(context.shape[0], context.shape[1], 3).to(x) mask_txt = torch.ones(context.shape[0], context.shape[1]).to(x.device, non_blocking=True).bool() return x, x_ids, txt, txt_ids, y, mask_x, mask_txt, x_seq_length def unpack(self, x: Tensor, cond: dict = None, x_seq_length: list = None) -> Tensor: x_list = [] image_shapes = cond["x_shapes"] for u, shape, seq_length in zip(x, image_shapes, x_seq_length): height, width = shape h, w = math.ceil(height / 2), math.ceil(width / 2) u = rearrange( u[seq_length-h*w:seq_length, ...], "(h w) (c ph pw) -> (h ph w pw) c", h=h, w=w, ph=2, pw=2, ) x_list.append(u) x = pad_sequence(tuple(x_list), batch_first=True).permute(0, 2, 1) return x def forward( self, x: Tensor, t: Tensor, cond: dict = {}, guidance: Tensor | None = None, gc_seg: int = 0, **kwargs ) -> Tensor: x, x_ids, txt, txt_ids, y, mask_x, mask_txt, seq_length_list = self.prepare_input(x, cond) # running on sequences img vec = self.time_in(timestep_embedding(t, 256)) if self.guidance_embed and guidance[-1] >= 0: 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) ids = torch.cat((txt_ids, x_ids), dim=1) pe = self.pe_embedder(ids) mask_aside = torch.cat((mask_txt, mask_x), dim=1) mask = mask_aside[:, None, :] * mask_aside[:, :, None] kwargs = dict( vec=vec, pe=pe, mask=mask, txt_length = txt.shape[1], ) x = torch.cat((txt, x), 1) if self.use_grad_checkpoint and gc_seg >= 0: x = checkpoint_sequential( functions=[partial(block, **kwargs) for block in self.double_blocks], segments=gc_seg if gc_seg > 0 else len(self.double_blocks), input=x, use_reentrant=False ) else: for block in self.double_blocks: x = block(x, **kwargs) kwargs = dict( vec=vec, pe=pe, mask=mask, ) if self.use_grad_checkpoint and gc_seg >= 0: x = checkpoint_sequential( functions=[partial(block, **kwargs) for block in self.single_blocks], segments=gc_seg if gc_seg > 0 else len(self.single_blocks), input=x, use_reentrant=False ) else: for block in self.single_blocks: x = block(x, **kwargs) x = x[:, txt.shape[1]:, ...] x = self.final_layer(x, vec) # (N, T, patch_size ** 2 * out_channels) 6 64 64 x = self.unpack(x, cond, seq_length_list) return x @staticmethod def get_config_template(): return dict_to_yaml('MODEL', __class__.__name__, FluxMR.para_dict, set_name=True) @BACKBONES.register_class() class FluxMRACEPlus(FluxMR): def __init__(self, cfg, logger = None): super().__init__(cfg, logger) def prepare_input(self, x, cond): context, y = cond["context"], cond["y"] batch_frames, batch_frames_ids = [], [] for ix, shape, imask, ie, ie_mask in zip(x, cond['x_shapes'], cond['x_mask'], cond['edit'], cond['edit_mask']): # unpack image from sequence ix = ix[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1]) imask = torch.ones_like( ix[[0], :, :]) if imask is None else imask.squeeze(0) if len(ie) > 0: ie = [iie.squeeze(0) for iie in ie] ie_mask = [ torch.ones( (ix.shape[0] * 4, ix.shape[1], ix.shape[2])) if iime is None else iime.squeeze(0) for iime in ie_mask ] ie = torch.cat(ie, dim=-1) ie_mask = torch.cat(ie_mask, dim=-1) else: ie, ie_mask = torch.zeros_like(ix).to(x), torch.ones_like( imask).to(x), ix = torch.cat([ix, ie, ie_mask], dim=0) c, h, w = ix.shape ix = rearrange(ix, 'c (h ph) (w pw) -> (h w) (c ph pw)', ph=2, pw=2) ix_id = torch.zeros(h // 2, w // 2, 3) ix_id[..., 1] = ix_id[..., 1] + torch.arange(h // 2)[:, None] ix_id[..., 2] = ix_id[..., 2] + torch.arange(w // 2)[None, :] ix_id = rearrange(ix_id, 'h w c -> (h w) c') batch_frames.append([ix]) batch_frames_ids.append([ix_id]) x_list, x_id_list, mask_x_list, x_seq_length = [], [], [], [] for frames, frame_ids in zip(batch_frames, batch_frames_ids): proj_frames = [] for idx, one_frame in enumerate(frames): one_frame = self.img_in(one_frame) proj_frames.append(one_frame) ix = torch.cat(proj_frames, dim=0) if_id = torch.cat(frame_ids, dim=0) x_list.append(ix) x_id_list.append(if_id) mask_x_list.append(torch.ones(ix.shape[0]).to(ix.device, non_blocking=True).bool()) x_seq_length.append(ix.shape[0]) # if len(x_list) < 1: import pdb;pdb.set_trace() x = pad_sequence(tuple(x_list), batch_first=True) x_ids = pad_sequence(tuple(x_id_list), batch_first=True).to(x) # [b,pad_seq,2] pad (0.,0.) at dim2 mask_x = pad_sequence(tuple(mask_x_list), batch_first=True) if isinstance(context, list): txt_list, mask_txt_list, y_list = [], [], [] for sample_id, (ctx, yy) in enumerate(zip(context, y)): txt_list.append(self.txt_in(ctx.to(x))) mask_txt_list.append(torch.ones(txt_list[-1].shape[0]).to(ctx.device, non_blocking=True).bool()) y_list.append(yy.to(x)) txt = pad_sequence(tuple(txt_list), batch_first=True) txt_ids = torch.zeros(txt.shape[0], txt.shape[1], 3).to(x) mask_txt = pad_sequence(tuple(mask_txt_list), batch_first=True) y = torch.cat(y_list, dim=0) assert y.ndim == 2 and txt.ndim == 3 else: txt = self.txt_in(context) txt_ids = torch.zeros(context.shape[0], context.shape[1], 3).to(x) mask_txt = torch.ones(context.shape[0], context.shape[1]).to(x.device, non_blocking=True).bool() return x, x_ids, txt, txt_ids, y, mask_x, mask_txt, x_seq_length @staticmethod def get_config_template(): return dict_to_yaml('MODEL', __class__.__name__, FluxMRACEPlus.para_dict, set_name=True) @BACKBONES.register_class() class FluxMRModiACEPlus(FluxMR): def __init__(self, cfg, logger = None): super().__init__(cfg, logger) def prepare_input(self, x, cond): context, y = cond["context"], cond["y"] batch_frames, batch_frames_ids = [], [] for ix, shape, imask, ie, im, ie_mask in zip(x, cond['x_shapes'], cond['x_mask'], cond['edit'], cond['modify'], cond['edit_mask']): # unpack image from sequence ix = ix[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1]) imask = torch.ones_like( ix[[0], :, :]) if imask is None else imask.squeeze(0) if len(ie) > 0: ie = [iie.squeeze(0) for iie in ie] im = [iim.squeeze(0) for iim in im] ie_mask = [ torch.ones( (ix.shape[0] * 4, ix.shape[1], ix.shape[2])) if iime is None else iime.squeeze(0) for iime in ie_mask ] im = torch.cat(im, dim=-1) ie = torch.cat(ie, dim=-1) ie_mask = torch.cat(ie_mask, dim=-1) else: ie, im, ie_mask = torch.zeros_like(ix).to(x), torch.zeros_like(ix).to(x), torch.ones_like( imask).to(x), ix = torch.cat([ix, ie, im, ie_mask], dim=0) c, h, w = ix.shape ix = rearrange(ix, 'c (h ph) (w pw) -> (h w) (c ph pw)', ph=2, pw=2) ix_id = torch.zeros(h // 2, w // 2, 3) ix_id[..., 1] = ix_id[..., 1] + torch.arange(h // 2)[:, None] ix_id[..., 2] = ix_id[..., 2] + torch.arange(w // 2)[None, :] ix_id = rearrange(ix_id, 'h w c -> (h w) c') batch_frames.append([ix]) batch_frames_ids.append([ix_id]) x_list, x_id_list, mask_x_list, x_seq_length = [], [], [], [] for frames, frame_ids in zip(batch_frames, batch_frames_ids): proj_frames = [] for idx, one_frame in enumerate(frames): one_frame = self.img_in(one_frame) proj_frames.append(one_frame) ix = torch.cat(proj_frames, dim=0) if_id = torch.cat(frame_ids, dim=0) x_list.append(ix) x_id_list.append(if_id) mask_x_list.append(torch.ones(ix.shape[0]).to(ix.device, non_blocking=True).bool()) x_seq_length.append(ix.shape[0]) # if len(x_list) < 1: import pdb;pdb.set_trace() x = pad_sequence(tuple(x_list), batch_first=True) x_ids = pad_sequence(tuple(x_id_list), batch_first=True).to(x) # [b,pad_seq,2] pad (0.,0.) at dim2 mask_x = pad_sequence(tuple(mask_x_list), batch_first=True) if isinstance(context, list): txt_list, mask_txt_list, y_list = [], [], [] for sample_id, (ctx, yy) in enumerate(zip(context, y)): txt_list.append(self.txt_in(ctx.to(x))) mask_txt_list.append(torch.ones(txt_list[-1].shape[0]).to(ctx.device, non_blocking=True).bool()) y_list.append(yy.to(x)) txt = pad_sequence(tuple(txt_list), batch_first=True) txt_ids = torch.zeros(txt.shape[0], txt.shape[1], 3).to(x) mask_txt = pad_sequence(tuple(mask_txt_list), batch_first=True) y = torch.cat(y_list, dim=0) assert y.ndim == 2 and txt.ndim == 3 else: txt = self.txt_in(context) txt_ids = torch.zeros(context.shape[0], context.shape[1], 3).to(x) mask_txt = torch.ones(context.shape[0], context.shape[1]).to(x.device, non_blocking=True).bool() return x, x_ids, txt, txt_ids, y, mask_x, mask_txt, x_seq_length @staticmethod def get_config_template(): return dict_to_yaml('MODEL', __class__.__name__, FluxMRACEPlus.para_dict, set_name=True)