diffsynth 1.0.0__py3-none-any.whl

diffsynth/models/svd_vae_encoder.py

@@ -0,0 +1,139 @@
+from .sd_vae_encoder import SDVAEEncoderStateDictConverter, SDVAEEncoder
+
+
+class SVDVAEEncoder(SDVAEEncoder):
+    def __init__(self):
+        super().__init__()
+        self.scaling_factor = 0.13025
+    
+    @staticmethod
+    def state_dict_converter():
+        return SVDVAEEncoderStateDictConverter()
+
+
+class SVDVAEEncoderStateDictConverter(SDVAEEncoderStateDictConverter):
+    def __init__(self):
+        super().__init__()
+
+    def from_diffusers(self, state_dict):
+        return super().from_diffusers(state_dict)
+    
+    def from_civitai(self, state_dict):
+        rename_dict = {
+            "conditioner.embedders.3.encoder.encoder.conv_in.bias": "conv_in.bias",
+            "conditioner.embedders.3.encoder.encoder.conv_in.weight": "conv_in.weight",
+            "conditioner.embedders.3.encoder.encoder.conv_out.bias": "conv_out.bias",
+            "conditioner.embedders.3.encoder.encoder.conv_out.weight": "conv_out.weight",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.0.conv1.bias": "blocks.0.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.0.conv1.weight": "blocks.0.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.0.conv2.bias": "blocks.0.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.0.conv2.weight": "blocks.0.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.0.norm1.bias": "blocks.0.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.0.norm1.weight": "blocks.0.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.0.norm2.bias": "blocks.0.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.0.norm2.weight": "blocks.0.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.1.conv1.bias": "blocks.1.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.1.conv1.weight": "blocks.1.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.1.conv2.bias": "blocks.1.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.1.conv2.weight": "blocks.1.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.1.norm1.bias": "blocks.1.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.1.norm1.weight": "blocks.1.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.1.norm2.bias": "blocks.1.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.0.block.1.norm2.weight": "blocks.1.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.0.downsample.conv.bias": "blocks.2.conv.bias",
+            "conditioner.embedders.3.encoder.encoder.down.0.downsample.conv.weight": "blocks.2.conv.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.conv1.bias": "blocks.3.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.conv1.weight": "blocks.3.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.conv2.bias": "blocks.3.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.conv2.weight": "blocks.3.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.nin_shortcut.bias": "blocks.3.conv_shortcut.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.nin_shortcut.weight": "blocks.3.conv_shortcut.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.norm1.bias": "blocks.3.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.norm1.weight": "blocks.3.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.norm2.bias": "blocks.3.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.0.norm2.weight": "blocks.3.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.1.conv1.bias": "blocks.4.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.1.conv1.weight": "blocks.4.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.1.conv2.bias": "blocks.4.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.1.conv2.weight": "blocks.4.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.1.norm1.bias": "blocks.4.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.1.norm1.weight": "blocks.4.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.1.norm2.bias": "blocks.4.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.block.1.norm2.weight": "blocks.4.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.1.downsample.conv.bias": "blocks.5.conv.bias",
+            "conditioner.embedders.3.encoder.encoder.down.1.downsample.conv.weight": "blocks.5.conv.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.conv1.bias": "blocks.6.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.conv1.weight": "blocks.6.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.conv2.bias": "blocks.6.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.conv2.weight": "blocks.6.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.nin_shortcut.bias": "blocks.6.conv_shortcut.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.nin_shortcut.weight": "blocks.6.conv_shortcut.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.norm1.bias": "blocks.6.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.norm1.weight": "blocks.6.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.norm2.bias": "blocks.6.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.0.norm2.weight": "blocks.6.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.1.conv1.bias": "blocks.7.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.1.conv1.weight": "blocks.7.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.1.conv2.bias": "blocks.7.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.1.conv2.weight": "blocks.7.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.1.norm1.bias": "blocks.7.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.1.norm1.weight": "blocks.7.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.1.norm2.bias": "blocks.7.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.block.1.norm2.weight": "blocks.7.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.2.downsample.conv.bias": "blocks.8.conv.bias",
+            "conditioner.embedders.3.encoder.encoder.down.2.downsample.conv.weight": "blocks.8.conv.weight",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.0.conv1.bias": "blocks.9.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.0.conv1.weight": "blocks.9.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.0.conv2.bias": "blocks.9.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.0.conv2.weight": "blocks.9.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.0.norm1.bias": "blocks.9.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.0.norm1.weight": "blocks.9.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.0.norm2.bias": "blocks.9.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.0.norm2.weight": "blocks.9.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.1.conv1.bias": "blocks.10.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.1.conv1.weight": "blocks.10.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.1.conv2.bias": "blocks.10.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.1.conv2.weight": "blocks.10.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.1.norm1.bias": "blocks.10.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.1.norm1.weight": "blocks.10.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.1.norm2.bias": "blocks.10.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.down.3.block.1.norm2.weight": "blocks.10.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.k.bias": "blocks.12.transformer_blocks.0.to_k.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.k.weight": "blocks.12.transformer_blocks.0.to_k.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.norm.bias": "blocks.12.norm.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.norm.weight": "blocks.12.norm.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.proj_out.bias": "blocks.12.transformer_blocks.0.to_out.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.proj_out.weight": "blocks.12.transformer_blocks.0.to_out.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.q.bias": "blocks.12.transformer_blocks.0.to_q.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.q.weight": "blocks.12.transformer_blocks.0.to_q.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.v.bias": "blocks.12.transformer_blocks.0.to_v.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.attn_1.v.weight": "blocks.12.transformer_blocks.0.to_v.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.block_1.conv1.bias": "blocks.11.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.block_1.conv1.weight": "blocks.11.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.block_1.conv2.bias": "blocks.11.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.block_1.conv2.weight": "blocks.11.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.block_1.norm1.bias": "blocks.11.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.block_1.norm1.weight": "blocks.11.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.block_1.norm2.bias": "blocks.11.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.block_1.norm2.weight": "blocks.11.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.block_2.conv1.bias": "blocks.13.conv1.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.block_2.conv1.weight": "blocks.13.conv1.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.block_2.conv2.bias": "blocks.13.conv2.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.block_2.conv2.weight": "blocks.13.conv2.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.block_2.norm1.bias": "blocks.13.norm1.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.block_2.norm1.weight": "blocks.13.norm1.weight",
+            "conditioner.embedders.3.encoder.encoder.mid.block_2.norm2.bias": "blocks.13.norm2.bias",
+            "conditioner.embedders.3.encoder.encoder.mid.block_2.norm2.weight": "blocks.13.norm2.weight",
+            "conditioner.embedders.3.encoder.encoder.norm_out.bias": "conv_norm_out.bias",
+            "conditioner.embedders.3.encoder.encoder.norm_out.weight": "conv_norm_out.weight",
+            "conditioner.embedders.3.encoder.quant_conv.bias": "quant_conv.bias",
+            "conditioner.embedders.3.encoder.quant_conv.weight": "quant_conv.weight",
+        }
+        state_dict_ = {}
+        for name in state_dict:
+            if name in rename_dict:
+                param = state_dict[name]
+                if "transformer_blocks" in rename_dict[name]:
+                    param = param.squeeze()
+                state_dict_[rename_dict[name]] = param
+        return state_dict_

diffsynth/models/tiler.py

@@ -0,0 +1,106 @@
+import torch
+from einops import rearrange, repeat
+
+
+class TileWorker:
+    def __init__(self):
+        pass
+
+
+    def mask(self, height, width, border_width):
+        # Create a mask with shape (height, width).
+        # The centre area is filled with 1, and the border line is filled with values in range (0, 1].
+        x = torch.arange(height).repeat(width, 1).T
+        y = torch.arange(width).repeat(height, 1)
+        mask = torch.stack([x + 1, height - x, y + 1, width - y]).min(dim=0).values
+        mask = (mask / border_width).clip(0, 1)
+        return mask
+
+
+    def tile(self, model_input, tile_size, tile_stride, tile_device, tile_dtype):
+        # Convert a tensor (b, c, h, w) to (b, c, tile_size, tile_size, tile_num)
+        batch_size, channel, _, _ = model_input.shape
+        model_input = model_input.to(device=tile_device, dtype=tile_dtype)
+        unfold_operator = torch.nn.Unfold(
+            kernel_size=(tile_size, tile_size),
+            stride=(tile_stride, tile_stride)
+        )
+        model_input = unfold_operator(model_input)
+        model_input = model_input.view((batch_size, channel, tile_size, tile_size, -1))
+
+        return model_input
+
+
+    def tiled_inference(self, forward_fn, model_input, tile_batch_size, inference_device, inference_dtype, tile_device, tile_dtype):
+        # Call y=forward_fn(x) for each tile
+        tile_num = model_input.shape[-1]
+        model_output_stack = []
+
+        for tile_id in range(0, tile_num, tile_batch_size):
+
+            # process input
+            tile_id_ = min(tile_id + tile_batch_size, tile_num)
+            x = model_input[:, :, :, :, tile_id: tile_id_]
+            x = x.to(device=inference_device, dtype=inference_dtype)
+            x = rearrange(x, "b c h w n -> (n b) c h w")
+
+            # process output
+            y = forward_fn(x)
+            y = rearrange(y, "(n b) c h w -> b c h w n", n=tile_id_-tile_id)
+            y = y.to(device=tile_device, dtype=tile_dtype)
+            model_output_stack.append(y)
+
+        model_output = torch.concat(model_output_stack, dim=-1)
+        return model_output
+
+
+    def io_scale(self, model_output, tile_size):
+        # Determine the size modification happend in forward_fn
+        # We only consider the same scale on height and width.
+        io_scale = model_output.shape[2] / tile_size
+        return io_scale
+    
+
+    def untile(self, model_output, height, width, tile_size, tile_stride, border_width, tile_device, tile_dtype):
+        # The reversed function of tile
+        mask = self.mask(tile_size, tile_size, border_width)
+        mask = mask.to(device=tile_device, dtype=tile_dtype)
+        mask = rearrange(mask, "h w -> 1 1 h w 1")
+        model_output = model_output * mask
+
+        fold_operator = torch.nn.Fold(
+            output_size=(height, width),
+            kernel_size=(tile_size, tile_size),
+            stride=(tile_stride, tile_stride)
+        )
+        mask = repeat(mask[0, 0, :, :, 0], "h w -> 1 (h w) n", n=model_output.shape[-1])
+        model_output = rearrange(model_output, "b c h w n -> b (c h w) n")
+        model_output = fold_operator(model_output) / fold_operator(mask)
+
+        return model_output
+
+
+    def tiled_forward(self, forward_fn, model_input, tile_size, tile_stride, tile_batch_size=1, tile_device="cpu", tile_dtype=torch.float32, border_width=None):
+        # Prepare
+        inference_device, inference_dtype = model_input.device, model_input.dtype
+        height, width = model_input.shape[2], model_input.shape[3]
+        border_width = int(tile_stride*0.5) if border_width is None else border_width
+
+        # tile
+        model_input = self.tile(model_input, tile_size, tile_stride, tile_device, tile_dtype)
+
+        # inference
+        model_output = self.tiled_inference(forward_fn, model_input, tile_batch_size, inference_device, inference_dtype, tile_device, tile_dtype)
+
+        # resize
+        io_scale = self.io_scale(model_output, tile_size)
+        height, width = int(height*io_scale), int(width*io_scale)
+        tile_size, tile_stride = int(tile_size*io_scale), int(tile_stride*io_scale)
+        border_width = int(border_width*io_scale)
+
+        # untile
+        model_output = self.untile(model_output, height, width, tile_size, tile_stride, border_width, tile_device, tile_dtype)
+        
+        # Done!
+        model_output = model_output.to(device=inference_device, dtype=inference_dtype)
+        return model_output

diffsynth/pipelines/__init__.py

@@ -0,0 +1,9 @@
+from .sd_image import SDImagePipeline
+from .sd_video import SDVideoPipeline
+from .sdxl_image import SDXLImagePipeline
+from .sdxl_video import SDXLVideoPipeline
+from .sd3_image import SD3ImagePipeline
+from .hunyuan_image import HunyuanDiTImagePipeline
+from .svd_video import SVDVideoPipeline
+from .pipeline_runner import SDVideoPipelineRunner
+KolorsImagePipeline = SDXLImagePipeline

diffsynth/pipelines/base.py

@@ -0,0 +1,34 @@
+import torch
+import numpy as np
+from PIL import Image
+
+
+
+class BasePipeline(torch.nn.Module):
+
+    def __init__(self, device="cuda", torch_dtype=torch.float16):
+        super().__init__()
+        self.device = device
+        self.torch_dtype = torch_dtype
+
+
+    def preprocess_image(self, image):
+        image = torch.Tensor(np.array(image, dtype=np.float32) * (2 / 255) - 1).permute(2, 0, 1).unsqueeze(0)
+        return image
+    
+
+    def preprocess_images(self, images):
+        return [self.preprocess_image(image) for image in images]
+    
+
+    def vae_output_to_image(self, vae_output):
+        image = vae_output[0].cpu().permute(1, 2, 0).numpy()
+        image = Image.fromarray(((image / 2 + 0.5).clip(0, 1) * 255).astype("uint8"))
+        return image
+    
+
+    def vae_output_to_video(self, vae_output):
+        video = vae_output.cpu().permute(1, 2, 0).numpy()
+        video = [Image.fromarray(((image / 2 + 0.5).clip(0, 1) * 255).astype("uint8")) for image in video]
+        return video
+    

diffsynth/pipelines/dancer.py

@@ -0,0 +1,178 @@
+import torch
+from ..models import SDUNet, SDMotionModel, SDXLUNet, SDXLMotionModel
+from ..models.sd_unet import PushBlock, PopBlock
+from ..controlnets import MultiControlNetManager
+
+
+def lets_dance(
+    unet: SDUNet,
+    motion_modules: SDMotionModel = None,
+    controlnet: MultiControlNetManager = None,
+    sample = None,
+    timestep = None,
+    encoder_hidden_states = None,
+    ipadapter_kwargs_list = {},
+    controlnet_frames = None,
+    unet_batch_size = 1,
+    controlnet_batch_size = 1,
+    cross_frame_attention = False,
+    tiled=False,
+    tile_size=64,
+    tile_stride=32,
+    device = "cuda",
+    vram_limit_level = 0,
+):
+    # 0. Text embedding alignment (only for video processing)
+    if encoder_hidden_states.shape[0] != sample.shape[0]:
+        encoder_hidden_states = encoder_hidden_states.repeat(sample.shape[0], 1, 1, 1)
+
+    # 1. ControlNet
+    #     This part will be repeated on overlapping frames if animatediff_batch_size > animatediff_stride.
+    #     I leave it here because I intend to do something interesting on the ControlNets.
+    controlnet_insert_block_id = 30
+    if controlnet is not None and controlnet_frames is not None:
+        res_stacks = []
+        # process controlnet frames with batch
+        for batch_id in range(0, sample.shape[0], controlnet_batch_size):
+            batch_id_ = min(batch_id + controlnet_batch_size, sample.shape[0])
+            res_stack = controlnet(
+                sample[batch_id: batch_id_],
+                timestep,
+                encoder_hidden_states[batch_id: batch_id_],
+                controlnet_frames[:, batch_id: batch_id_],
+                tiled=tiled, tile_size=tile_size, tile_stride=tile_stride
+            )
+            if vram_limit_level >= 1:
+                res_stack = [res.cpu() for res in res_stack]
+            res_stacks.append(res_stack)
+        # concat the residual
+        additional_res_stack = []
+        for i in range(len(res_stacks[0])):
+            res = torch.concat([res_stack[i] for res_stack in res_stacks], dim=0)
+            additional_res_stack.append(res)
+    else:
+        additional_res_stack = None
+
+    # 2. time
+    time_emb = unet.time_proj(timestep).to(sample.dtype)
+    time_emb = unet.time_embedding(time_emb)
+
+    # 3. pre-process
+    height, width = sample.shape[2], sample.shape[3]
+    hidden_states = unet.conv_in(sample)
+    text_emb = encoder_hidden_states
+    res_stack = [hidden_states.cpu() if vram_limit_level>=1 else hidden_states]
+
+    # 4. blocks
+    for block_id, block in enumerate(unet.blocks):
+        # 4.1 UNet
+        if isinstance(block, PushBlock):
+            hidden_states, time_emb, text_emb, res_stack = block(hidden_states, time_emb, text_emb, res_stack)
+            if vram_limit_level>=1:
+                res_stack[-1] = res_stack[-1].cpu()
+        elif isinstance(block, PopBlock):
+            if vram_limit_level>=1:
+                res_stack[-1] = res_stack[-1].to(device)
+            hidden_states, time_emb, text_emb, res_stack = block(hidden_states, time_emb, text_emb, res_stack)
+        else:
+            hidden_states_input = hidden_states
+            hidden_states_output = []
+            for batch_id in range(0, sample.shape[0], unet_batch_size):
+                batch_id_ = min(batch_id + unet_batch_size, sample.shape[0])
+                hidden_states, _, _, _ = block(
+                    hidden_states_input[batch_id: batch_id_],
+                    time_emb,
+                    text_emb[batch_id: batch_id_],
+                    res_stack,
+                    cross_frame_attention=cross_frame_attention,
+                    ipadapter_kwargs_list=ipadapter_kwargs_list.get(block_id, {}),
+                    tiled=tiled, tile_size=tile_size, tile_stride=tile_stride
+                )
+                hidden_states_output.append(hidden_states)
+            hidden_states = torch.concat(hidden_states_output, dim=0)
+        # 4.2 AnimateDiff
+        if motion_modules is not None:
+            if block_id in motion_modules.call_block_id:
+                motion_module_id = motion_modules.call_block_id[block_id]
+                hidden_states, time_emb, text_emb, res_stack = motion_modules.motion_modules[motion_module_id](
+                    hidden_states, time_emb, text_emb, res_stack,
+                    batch_size=1
+                )
+        # 4.3 ControlNet
+        if block_id == controlnet_insert_block_id and additional_res_stack is not None:
+            hidden_states += additional_res_stack.pop().to(device)
+            if vram_limit_level>=1:
+                res_stack = [(res.to(device) + additional_res.to(device)).cpu() for res, additional_res in zip(res_stack, additional_res_stack)]
+            else:
+                res_stack = [res + additional_res for res, additional_res in zip(res_stack, additional_res_stack)]
+    
+    # 5. output
+    hidden_states = unet.conv_norm_out(hidden_states)
+    hidden_states = unet.conv_act(hidden_states)
+    hidden_states = unet.conv_out(hidden_states)
+
+    return hidden_states
+
+
+
+
+def lets_dance_xl(
+    unet: SDXLUNet,
+    motion_modules: SDXLMotionModel = None,
+    controlnet: MultiControlNetManager = None,
+    sample = None,
+    add_time_id = None,
+    add_text_embeds = None,
+    timestep = None,
+    encoder_hidden_states = None,
+    ipadapter_kwargs_list = {},
+    controlnet_frames = None,
+    unet_batch_size = 1,
+    controlnet_batch_size = 1,
+    cross_frame_attention = False,
+    tiled=False,
+    tile_size=64,
+    tile_stride=32,
+    device = "cuda",
+    vram_limit_level = 0,
+):
+    # 2. time
+    t_emb = unet.time_proj(timestep).to(sample.dtype)
+    t_emb = unet.time_embedding(t_emb)
+
+    time_embeds = unet.add_time_proj(add_time_id)
+    time_embeds = time_embeds.reshape((add_text_embeds.shape[0], -1))
+    add_embeds = torch.concat([add_text_embeds, time_embeds], dim=-1)
+    add_embeds = add_embeds.to(sample.dtype)
+    add_embeds = unet.add_time_embedding(add_embeds)
+
+    time_emb = t_emb + add_embeds
+
+    # 3. pre-process
+    height, width = sample.shape[2], sample.shape[3]
+    hidden_states = unet.conv_in(sample)
+    text_emb = encoder_hidden_states if unet.text_intermediate_proj is None else unet.text_intermediate_proj(encoder_hidden_states)
+    res_stack = [hidden_states]
+
+    # 4. blocks
+    for block_id, block in enumerate(unet.blocks):
+        hidden_states, time_emb, text_emb, res_stack = block(
+            hidden_states, time_emb, text_emb, res_stack,
+            tiled=tiled, tile_size=tile_size, tile_stride=tile_stride,
+            ipadapter_kwargs_list=ipadapter_kwargs_list.get(block_id, {})
+        )
+        # 4.2 AnimateDiff
+        if motion_modules is not None:
+            if block_id in motion_modules.call_block_id:
+                motion_module_id = motion_modules.call_block_id[block_id]
+                hidden_states, time_emb, text_emb, res_stack = motion_modules.motion_modules[motion_module_id](
+                    hidden_states, time_emb, text_emb, res_stack,
+                    batch_size=1
+                )
+
+    # 5. output
+    hidden_states = unet.conv_norm_out(hidden_states)
+    hidden_states = unet.conv_act(hidden_states)
+    hidden_states = unet.conv_out(hidden_states)
+
+    return hidden_states