Init

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+outputs

README.md

@@ -11,3 +11,5 @@ license: cc-by-nc-sa-4.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+arxiv.org/abs/2408.04631

app.py

@@ -0,0 +1,502 @@
+import os
+from PIL import Image, ImageOps
+import random
+
+import cv2
+from diffusers import StableVideoDiffusionPipeline
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.attention_processor import XFormersAttnProcessor
+from diffusers.utils import export_to_gif
+import gradio as gr
+import numpy as np
+from safetensors import safe_open
+from segment_anything import build_sam, SamPredictor
+import spaces
+from tqdm import tqdm
+import torch
+
+from svd import (
+    UNetDragSpatioTemporalConditionModel,
+    AllToFirstXFormersAttnProcessor,
+)
+
+
+TITLE = '''Puppet-Master: Scaling Interactive Video Generation as a Motion Prior for Part-Level Dynamics'''
+DESCRIPTION = """
+<div>
+Try <a href='https://vgg-puppetmaster.github.io/'><b>Puppet-Master</b></a> yourself to animate your favorite objects in seconds!
+</div>
+<div>
+Please give us a 🌟 on <a href='https://github.com/RuiningLi/puppet-master'>Github</a> if you like our work!
+</div>
+"""
+INSTRUCTION = '''
+2 steps to get started:
+- Upload an image of a dynamic object.
+- Add one or more drags on the object to specify the part-level interactions.
+How to add drags:
+- To add a drag, first click on the starting point of the drag, then click on the ending point of the drag, on the Input Image (leftmost).
+- You can add up to 5 drags.
+- After every click, the drags will be visualized on the Image with Drags (second from left).
+- If the last drag is not completed (you specified the starting point but not the ending point), it will simply be ignored.
+- To retry, click the [x] button on the top-right corner of the input image to start over, even if you just want to try a different set of drags.
+- Have fun dragging!
+
+Then, you will be prompted to verify the object segmentation. Once you confirm that the segmentation is decent, the output image will be generated in seconds!
+
+Tips:
+- We found having classifier-free guidance weight ~5.0 works best.
+- Try changing the random seed to get different results.
+'''
+PREPROCESS_INSTRUCTION = '''
+Segmentation is needed if it is not already provided through an alpha channel in the input image.
+You don't need to tick this box if you have chosen one of the example images.
+If you have uploaded one of your own images, it is very likely that you will need to tick this box.
+You should verify that the preprocessed image is object-centric (i.e., clearly contains a single object) and has white background.
+'''
+
+
+def tensor2vid(video: torch.Tensor, processor: VaeImageProcessor, output_type: str = "np"):
+    batch_size = video.shape[0]
+    outputs = []
+    for batch_idx in range(batch_size):
+        batch_vid = video[batch_idx].permute(1, 0, 2, 3)
+        batch_output = processor.postprocess(batch_vid, output_type)
+
+        outputs.append(batch_output)
+
+    if output_type == "np":
+        outputs = np.stack(outputs)
+
+    elif output_type == "pt":
+        outputs = torch.stack(outputs)
+
+    elif not output_type == "pil":
+        raise ValueError(f"{output_type} does not exist. Please choose one of ['np', 'pt', 'pil']")
+
+    return outputs
+
+
+def center_and_square_image(pil_image_rgba, drags, scale_factor):
+    image = pil_image_rgba
+    alpha = np.array(image)[:, :, 3]  # Extract the alpha channel
+
+    foreground_coords = np.argwhere(alpha > 0)
+    y_min, x_min = foreground_coords.min(axis=0)
+    y_max, x_max = foreground_coords.max(axis=0)
+    cy, cx = (y_min + y_max) // 2, (x_min + x_max) // 2
+    crop_height, crop_width = y_max - y_min + 1, x_max - x_min + 1
+    side_length = int(max(crop_height, crop_width) * scale_factor)
+    padded_image = ImageOps.expand(
+        image, 
+        (side_length // 2, side_length // 2, side_length // 2, side_length // 2), 
+        fill=(255, 255, 255, 255)
+    )
+    left, top = cx, cy
+    new_drags = []
+    for d in drags:
+        x, y = d
+        new_x, new_y = (x + side_length // 2 - cx) / side_length, (y + side_length // 2 - cy) / side_length
+        new_drags.append((new_x, new_y))
+
+    # Crop or pad the image as needed to make it centered around (cx, cy)
+    image = padded_image.crop((left, top, left + side_length, top + side_length))
+    # Resize the image to 256x256
+    image = image.resize((256, 256), Image.Resampling.LANCZOS)
+    return image, new_drags
+
+
+def sam_init():
+    sam_checkpoint = os.path.join(os.path.dirname(__file__), "ckpts", "sam_vit_h_4b8939.pth")
+    predictor = SamPredictor(build_sam(checkpoint=sam_checkpoint).to("cuda"))
+    return predictor
+
+
+def model_init():
+    model_checkpoint = os.path.join(os.path.dirname(__file__), "ckpts", "model.safetensors")
+    state_dict = {}
+    with safe_open(model_checkpoint, framework="pt", device="cpu") as f:
+        for k in f.keys():
+            state_dict[k] = f.get_tensor(k)
+    model = UNetDragSpatioTemporalConditionModel(num_drags=5)
+    attn_processors_dict={
+        "down_blocks.0.attentions.0.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "down_blocks.0.attentions.0.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.0.attentions.1.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "down_blocks.0.attentions.1.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.1.attentions.0.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "down_blocks.1.attentions.0.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.1.attentions.1.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "down_blocks.1.attentions.1.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.2.attentions.0.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "down_blocks.2.attentions.0.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.2.attentions.1.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "down_blocks.2.attentions.1.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+
+        "down_blocks.0.attentions.0.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "down_blocks.0.attentions.0.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.0.attentions.1.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "down_blocks.0.attentions.1.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.1.attentions.0.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "down_blocks.1.attentions.0.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.1.attentions.1.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "down_blocks.1.attentions.1.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.2.attentions.0.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "down_blocks.2.attentions.0.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "down_blocks.2.attentions.1.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "down_blocks.2.attentions.1.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+
+        "up_blocks.1.attentions.0.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.1.attentions.1.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "up_blocks.1.attentions.1.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.1.attentions.2.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "up_blocks.1.attentions.2.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.2.attentions.0.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "up_blocks.2.attentions.0.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.2.attentions.1.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "up_blocks.2.attentions.1.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.2.attentions.2.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "up_blocks.2.attentions.2.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.3.attentions.0.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "up_blocks.3.attentions.0.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.3.attentions.1.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "up_blocks.3.attentions.1.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.3.attentions.2.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "up_blocks.3.attentions.2.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+
+        "up_blocks.1.attentions.0.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "up_blocks.1.attentions.0.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.1.attentions.1.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "up_blocks.1.attentions.1.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.1.attentions.2.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "up_blocks.1.attentions.2.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.2.attentions.0.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "up_blocks.2.attentions.0.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.2.attentions.1.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "up_blocks.2.attentions.1.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.2.attentions.2.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "up_blocks.2.attentions.2.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.3.attentions.0.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "up_blocks.3.attentions.0.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.3.attentions.1.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "up_blocks.3.attentions.1.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "up_blocks.3.attentions.2.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "up_blocks.3.attentions.2.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+
+        "mid_block.attentions.0.transformer_blocks.0.attn1.processor": AllToFirstXFormersAttnProcessor(),
+        "mid_block.attentions.0.transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+        "mid_block.attentions.0.temporal_transformer_blocks.0.attn1.processor": XFormersAttnProcessor(),
+        "mid_block.attentions.0.temporal_transformer_blocks.0.attn2.processor": XFormersAttnProcessor(),
+    }
+    model.set_attn_processor(attn_processors_dict)
+    model.load_state_dict(state_dict, strict=True)
+    return model.to("cuda")
+
+
+sam_predictor = sam_init()
+model = model_init()
+pipe = StableVideoDiffusionPipeline.from_pretrained(
+    "/scratch/shared/beegfs/ruining/projects/generative-models/stable-video-diffusion-img2vid",
+    torch_dtype=torch.float16, variant="fp16"
+)
+pipe.vae.to(dtype=torch.float16, device="cuda")
+pipe.image_encoder = pipe.image_encoder.to("cuda")
+
+
+@spaces.GPU(duration=10)
+def sam_segment(input_image, drags, foreground_points=None, scale_factor=2.2):
+    image = np.asarray(input_image)
+    sam_predictor.set_image(image)
+
+    with torch.no_grad():
+        masks_bbox, _, _ = sam_predictor.predict(
+            point_coords=foreground_points if foreground_points is not None else None,
+            point_labels=np.ones(len(foreground_points)) if foreground_points is not None else None,
+            multimask_output=True
+        )
+
+    out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
+    out_image[:, :, :3] = image
+    out_image[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255
+    torch.cuda.empty_cache()
+    out_image, new_drags = center_and_square_image(Image.fromarray(out_image, mode="RGBA"), drags, scale_factor)
+
+    return out_image, new_drags
+
+
+def get_point(img, sel_pix, evt: gr.SelectData):
+    sel_pix.append(evt.index)
+    points = []
+    img = np.array(img)
+    height = img.shape[0]
+    arrow_width_large = 7 * height // 256
+    arrow_width_small = 3 * height // 256
+    circle_size = 5 * height // 256
+
+    with_alpha = img.shape[2] == 4
+    for idx, point in enumerate(sel_pix):
+        if idx % 2 == 1:
+            cv2.circle(img, tuple(point), circle_size, (0, 0, 255, 255) if with_alpha else (0, 0, 255), -1)
+        else:
+            cv2.circle(img, tuple(point), circle_size, (255, 0, 0, 255) if with_alpha else (255, 0, 0), -1)
+        points.append(tuple(point))
+        if len(points) == 2:
+            cv2.arrowedLine(img, points[0], points[1], (0, 0, 0, 255) if with_alpha else (0, 0, 0), arrow_width_large)
+            cv2.arrowedLine(img, points[0], points[1], (255, 255, 0, 255) if with_alpha else (0, 0, 0), arrow_width_small)
+            points = []
+    return img if isinstance(img, np.ndarray) else np.array(img)
+
+
+def clear_drag():
+    return []
+
+
+def preprocess_image(img, chk_group, drags):
+    if img is None:
+        gr.Warning("No image is specified. Please specify an image before preprocessing.")
+        return None, drags
+
+    if drags is None or len(drags) == 0:
+        foreground_points = None
+    else:
+        foreground_points = np.array([drags[i] for i in range(0, len(drags), 2)])
+
+    if len(drags) == 0:
+        gr.Warning("No drags are specified. We recommend first specifying the drags before preprocessing.")
+
+    new_drags = drags
+    if "Preprocess with Segmentation" in chk_group:
+        img_np = np.array(img)
+        rgb_img = img_np[..., :3]
+        img, new_drags = sam_segment(
+            rgb_img,
+            drags,
+            foreground_points=foreground_points,
+        )
+    else:
+        new_drags = [(d[0] / img.width, d[1] / img.height) for d in drags]
+
+    img = np.array(img).astype(np.float32)
+    processed_img = img[..., :3] * img[..., 3:] / 255. + 255. * (1 - img[..., 3:] / 255.)
+    image_pil = Image.fromarray(processed_img.astype(np.uint8), mode="RGB")
+    processed_img = image_pil.resize((256, 256), Image.LANCZOS)
+    return processed_img, new_drags
+
+
+def sample_from_noise(model, scheduler, cond_latent, cond_embedding, drags,
+                      min_guidance=1.0, max_guidance=3.0, num_inference_steps=50):
+    model.eval()
+
+    scheduler.set_timesteps(num_inference_steps, device=cond_latent.device)
+    timesteps = scheduler.timesteps.to(cond_latent.device)
+
+    do_classifier_free_guidance = max_guidance > 1.0
+    latents = torch.randn((1, 14, 4, 32, 32)).to(cond_latent) * scheduler.init_noise_sigma
+    guidance_scale = torch.linspace(min_guidance, max_guidance, 14).unsqueeze(0).to(cond_latent)[..., None, None, None]
+
+    for i, t in tqdm(enumerate(timesteps)):
+        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+
+        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
+
+        with torch.no_grad():
+            noise_pred = model(
+                latent_model_input,
+                t,
+                image_latents=torch.cat([cond_latent, torch.zeros_like(cond_latent)]) if do_classifier_free_guidance else cond_latent,
+                encoder_hidden_states=torch.cat([cond_embedding, torch.zeros_like(cond_embedding)]) if do_classifier_free_guidance else cond_embedding,
+                added_time_ids=None,  # dummy
+                drags=torch.cat([drags, torch.zeros_like(drags)]) if do_classifier_free_guidance else drags,
+            )
+
+        if do_classifier_free_guidance:
+            noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+        
+        latents = scheduler.step(noise_pred, t, latents).prev_sample
+    
+    return latents
+
+
+@spaces.GPU(duration=40)
+def generate_image(img_cond, seed, cfg_scale, drags_list):
+    if img_cond is None:
+        gr.Warning("Please preprocess the image first.")
+        return None
+    
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+    img_cond_pil = Image.fromarray(img_cond)
+    img_cond_preprocessed = pipe.video_processor.preprocess(img_cond_pil, height=256, width=256)
+    img_cond_preprocessed = img_cond_preprocessed.to(device="cuda", dtype=torch.float16)
+    latent_dist = pipe.vae.encode(img_cond_preprocessed).latent_dist
+    embeddings = pipe._encode_image(img_cond_pil, device="cuda", num_videos_per_prompt=1, do_classifier_free_guidance=False)
+
+    drags = torch.zeros(14, 5, 4)
+    for i in range(0, len(drags_list), 2):
+        start_point, end_point = drags_list[i:i+2]
+        drag_idx = i // 2
+        drags[:, drag_idx, :2] = torch.Tensor(start_point)
+        drags[0, drag_idx, 2:] = torch.Tensor(start_point)
+        drags[-1, drag_idx, 2:] = torch.Tensor(end_point)
+
+        if drag_idx == 4:
+            break
+
+    frame_indices = torch.arange(1, 13).unsqueeze(-1).unsqueeze(-1)
+    t = frame_indices.float() / 13.0  # Normalize time to [0, 1]
+    drags[1:-1, :, 2:] = drags[0, :, 2:] * (1 - t) + drags[-1, :, 2:] * t
+    drags = drags[None].to(device="cuda")
+
+    batch = dict(
+        drags=drags,
+        cond_embedding=embeddings.to(dtype=torch.float32),
+        cond_latent=latent_dist.mean.to(dtype=torch.float32),
+    )
+
+    with torch.no_grad():
+        latents = sample_from_noise(
+            model,
+            pipe.scheduler,
+            **batch,
+            max_guidance=cfg_scale,
+            num_inference_steps=50,
+        )
+
+        frames = pipe.vae.decode(latents.flatten(0, 1).to(torch.float16) / 0.18215, num_frames=14).sample.float()
+        frames = tensor2vid(frames.view(-1, 14, 3, 256, 256).permute(0, 2, 1, 3, 4), pipe.video_processor, output_type="pil")[0]
+
+    # Add drags
+    frame_with_drag = np.ascontiguousarray(np.array(frames[0]))
+    for i in range(0, len(drags_list), 2):
+        drag_idx = i // 2
+        start_point, end_point = drags_list[i:i+2]
+        start_point = (int(start_point[0] * 256), int(start_point[1] * 256))
+        end_point = (int(end_point[0] * 256), int(end_point[1] * 256))
+        frame_with_drag = cv2.arrowedLine(frame_with_drag, start_point, end_point, (0, 0, 0), 4)
+        frame_with_drag = cv2.arrowedLine(frame_with_drag, start_point, end_point, (255, 255, 0), 2)
+
+        if drag_idx == 4:
+            break
+
+    frames = [Image.fromarray(frame_with_drag)] * 5 + frames
+    save_dir = os.path.join(os.path.dirname(__file__), "outputs")
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    save_id = len(os.listdir(save_dir))
+    save_path = os.path.join(save_dir, f"{save_id:05d}.gif")
+    export_to_gif(frames, save_path)
+    return save_path
+
+
+with gr.Blocks(title=TITLE) as demo:
+    gr.Markdown("# " + DESCRIPTION)
+
+    with gr.Row():
+        gr.Markdown(INSTRUCTION)
+    
+    drags = gr.State(value=[])
+
+    with gr.Row(variant="panel"):
+        with gr.Column(scale=1):
+            input_image = gr.Image(
+                interactive=True,
+                type='pil',
+                image_mode="RGBA",
+                width=256,
+                show_label=True,
+                label="Input Image",
+            )
+
+            example_folder = os.path.join(os.path.dirname(__file__), "./example_images")
+            example_fns = [os.path.join(example_folder, example) for example in sorted(os.listdir(example_folder))]
+            gr.Examples(
+                examples=example_fns,
+                inputs=[input_image],
+                cache_examples=False,
+                label='Feel free to use one of our provided examples!',
+                examples_per_page=30
+            )
+
+            input_image.change(
+                fn=clear_drag,
+                outputs=[drags],
+            )
+
+        with gr.Column(scale=1):
+            drag_image = gr.Image(
+                type="numpy",
+                label="Image with Drags",
+                interactive=False,
+                width=256,
+                image_mode="RGB",
+            )
+
+            input_image.select(
+                fn=get_point,
+                inputs=[input_image, drags],
+                outputs=[drag_image],
+            )
+        
+        with gr.Column(scale=1):
+            processed_image = gr.Image(
+                type='numpy', 
+                label="Processed Image", 
+                interactive=False, 
+                width=256,
+                height=256,
+                image_mode='RGB',
+            )
+            processed_image_highres = gr.Image(type='pil', image_mode='RGB', visible=False)
+
+            with gr.Accordion('Advanced preprocessing options', open=True):
+                with gr.Row():
+                    with gr.Column():
+                        preprocess_chk_group = gr.CheckboxGroup(
+                            ['Preprocess with Segmentation'], 
+                            label='Segment',
+                            info=PREPROCESS_INSTRUCTION
+                        )
+            
+            preprocess_button = gr.Button(
+                value="Preprocess Input Image",
+            )
+            preprocess_button.click(
+                fn=preprocess_image,
+                inputs=[input_image, preprocess_chk_group, drags],
+                outputs=[processed_image, drags],
+                queue=True,
+            )
+
+        with gr.Column(scale=1):
+            generated_gif = gr.Image(
+                type="filepath",
+                label="Generated GIF",
+                interactive=False,
+                height=256,
+                width=256,
+                image_mode="RGB",
+            )
+
+            with gr.Accordion('Advanced generation options', open=True):
+                with gr.Row():
+                    with gr.Column():
+                        seed = gr.Slider(label="seed", value=0, minimum=0, maximum=10000, step=1, randomize=False)
+                        cfg_scale = gr.Slider(
+                            label="classifier-free guidance weight",
+                            value=5, minimum=1, maximum=10, step=0.1
+                        )
+
+            generate_button = gr.Button(
+                value="Generate Image",
+            )
+            generate_button.click(
+                fn=generate_image,
+                inputs=[processed_image, seed, cfg_scale, drags],
+                outputs=[generated_gif],
+            )
+
+    demo.launch(share=True)

ckpts/model.safetensors

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

ckpts/sam_vit_h_4b8939.pth

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

network_utils.py

@@ -0,0 +1,104 @@
+from typing import Tuple
+
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+
+class DragEmbedding(nn.Module):
+    def __init__(
+        self,
+        conditioning_embedding_channels: int,  # out channel
+        conditioning_channels: int = 3,
+        block_out_channels: Tuple[int, ...] = (16, 32, 96),
+    ):
+        super().__init__()
+
+        self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1)
+
+        self.blocks = nn.ModuleList([])
+
+        for i in range(len(block_out_channels) - 1):
+            channel_in = block_out_channels[i]
+            channel_out = block_out_channels[i + 1]
+            self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1))
+            self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1))
+
+        self.conv_out = zero_module(
+            nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)
+        )
+
+    def forward(self, conditioning):
+        conditioning_ndims = len(conditioning.shape)
+        if conditioning_ndims == 5:
+            batch_size, num_frames, num_channels, h, w = conditioning.shape
+            conditioning = conditioning.flatten(0, 1)
+
+        embedding = self.conv_in(conditioning)
+        embedding = F.silu(embedding)
+
+        for block in self.blocks:
+            embedding = block(embedding)
+            embedding = F.silu(embedding)
+
+        embedding = self.conv_out(embedding)
+        if conditioning_ndims == 5:
+            embedding = embedding.view(batch_size, num_frames, *embedding.shape[1:])
+
+        return embedding
+    
+
+def zero_module(module):
+    for p in module.parameters():
+        nn.init.zeros_(p)
+    return module
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb

package-lock.json

@@ -0,0 +1,6 @@
+{
+  "name": "puppet-master",
+  "lockfileVersion": 2,
+  "requires": true,
+  "packages": {}
+}

requirements.txt

@@ -0,0 +1,154 @@
+accelerate==0.31.0
+aiofiles==23.2.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+aiostream==0.5.2
+altair==5.3.0
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.4.0
+async-timeout==4.0.3
+attrs==23.2.0
+beautifulsoup4==4.12.3
+Brotli==1.1.0
+certifi==2024.6.2
+charset-normalizer==3.3.2
+click==8.1.7
+cloudpickle==3.0.0
+contourpy==1.2.1
+cycler==0.12.1
+diffusers @ git+https://github.com/huggingface/diffusers@8e1b7a084addc4711b8d9be2738441dfad680ce0
+dnspython==2.6.1
+docker-pycreds==0.4.0
+email_validator==2.2.0
+exceptiongroup==1.2.1
+fastapi==0.111.0
+fastapi-cli==0.0.4
+ffmpy==0.3.2
+filelock==3.13.1
+fonttools==4.53.0
+frozenlist==1.4.1
+fsspec==2024.2.0
+gdown==5.2.0
+gitdb==4.0.11
+GitPython==3.1.43
+gradio
+gradio_client
+grpclib==0.4.7
+h11==0.14.0
+h2==4.1.0
+h5py==3.11.0
+hpack==4.0.0
+httpcore==1.0.5
+httptools==0.6.1
+httpx==0.27.0
+huggingface-hub==0.23.4
+hydra-core==1.3.1
+hyperframe==6.0.1
+idna==3.7
+imageio==2.34.1
+importlib_metadata==7.1.0
+importlib_resources==6.4.0
+Jinja2==3.1.3
+jsonschema==4.22.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+lazy_loader==0.4
+lightning-utilities==0.11.2
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.0
+mdurl==0.1.2
+modal==0.63.2
+mpmath==1.3.0
+multidict==6.0.5
+mutagen==1.47.0
+networkx==3.2.1
+numpy==1.26.3
+nvidia-cublas-cu11==11.11.3.6
+nvidia-cuda-cupti-cu11==11.8.87
+nvidia-cuda-nvrtc-cu11==11.8.89
+nvidia-cuda-runtime-cu11==11.8.89
+nvidia-cudnn-cu11==8.7.0.84
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.3.0.86
+nvidia-cusolver-cu11==11.4.1.48
+nvidia-cusparse-cu11==11.7.5.86
+nvidia-nccl-cu11==2.20.5
+nvidia-nvtx-cu11==11.8.86
+omegaconf==2.3.0
+opencv-python==4.10.0.82
+orjson==3.10.5
+packaging==24.1
+pandas==2.2.2
+pillow==10.2.0
+platformdirs==4.2.2
+protobuf==4.25.3
+psutil==5.9.8
+pycryptodomex==3.20.0
+pydantic==2.7.4
+pydantic_core==2.18.4
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.1.2
+PySocks==1.7.1
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.9
+pytube==15.0.0
+pytube3==9.6.4
+pytz==2024.1
+PyYAML==6.0.1
+referencing==0.35.1
+regex==2024.5.15
+requests==2.32.3
+rich==13.7.1
+rpds-py==0.18.1
+ruff==0.5.0
+safetensors==0.4.3
+scikit-image==0.24.0
+scipy==1.13.1
+segment-anything==1.0
+semantic-version==2.10.0
+sentry-sdk==2.5.1
+setproctitle==1.3.3
+shellingham==1.5.4
+sigtools==4.0.1
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.1
+soupsieve==2.5
+spaces==0.29.3
+starlette==0.37.2
+submitit==1.5.1
+sympy==1.12
+synchronicity==0.6.7
+tifffile==2024.6.18
+tokenizers==0.19.1
+toml==0.10.2
+tomlkit==0.12.0
+toolz==0.12.1
+torch==2.3.0+cu118
+torchaudio==2.3.1+cu118
+torchmetrics==1.4.0.post0
+torchvision==0.18.1+cu118
+tqdm==4.66.4
+transformers==4.41.2
+triton==2.3.0
+typer==0.12.3
+types-certifi==2021.10.8.3
+types-toml==0.10.8.20240310
+typing_extensions==4.9.0
+tzdata==2024.1
+ujson==5.10.0
+urllib3==2.2.1
+uvicorn==0.30.1
+uvloop==0.19.0
+wandb==0.17.1
+watchfiles==0.22.0
+websockets==11.0.3
+xformers==0.0.26.post1+cu118
+yarl==1.9.4
+youtube-dl==2021.12.17
+yt-dlp==2024.5.27
+zipp==3.19.2

svd.py

@@ -0,0 +1,1316 @@
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple, Union, Any, Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import UNet2DConditionLoadersMixin
+from diffusers.utils import BaseOutput, logging
+from diffusers.utils.torch_utils import is_torch_version
+from diffusers.models.attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.unets.unet_3d_blocks import (
+    UNetMidBlockSpatioTemporal,
+    get_down_block as gdb, 
+    get_up_block as gub,
+)
+from diffusers.models.resnet import (
+    Downsample2D,
+    SpatioTemporalResBlock,
+    Upsample2D,
+)
+from diffusers.models.transformers.transformer_temporal import TransformerSpatioTemporalModel
+from diffusers.models.attention_processor import Attention
+from diffusers.utils import deprecate
+from diffusers.utils.import_utils import is_xformers_available
+
+from network_utils import DragEmbedding, get_2d_sincos_pos_embed
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+if is_xformers_available():
+    import xformers
+    import xformers.ops
+
+
+class AllToFirstXFormersAttnProcessor:
+    r"""
+    Processor for implementing memory efficient attention using xFormers.
+
+    Args:
+        attention_op (`Callable`, *optional*, defaults to `None`):
+            The base
+            [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to
+            use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best
+            operator.
+    """
+
+    def __init__(self, attention_op: Optional[Callable] = None):
+        self.attention_op = attention_op
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        temb: Optional[torch.FloatTensor] = None,
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+            deprecate("scale", "1.0.0", deprecation_message)
+
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, key_tokens, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        assert encoder_hidden_states is None
+        attention_mask = attn.prepare_attention_mask(attention_mask, key_tokens, batch_size)
+        if attention_mask is not None:
+            # expand our mask's singleton query_tokens dimension:
+            #   [batch*heads,            1, key_tokens] ->
+            #   [batch*heads, query_tokens, key_tokens]
+            # so that it can be added as a bias onto the attention scores that xformers computes:
+            #   [batch*heads, query_tokens, key_tokens]
+            # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
+            _, query_tokens, _ = hidden_states.shape
+            attention_mask = attention_mask.expand(-1, query_tokens, -1)
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states.view(-1, 14, *hidden_states.shape[1:])[:, 0])[:, None].expand(-1, 14, -1, -1).flatten(0, 1)
+        value = attn.to_v(hidden_states.view(-1, 14, *hidden_states.shape[1:])[:, 0])[:, None].expand(-1, 14, -1, -1).flatten(0, 1)
+
+        query = attn.head_to_batch_dim(query).contiguous()
+        key = attn.head_to_batch_dim(key).contiguous()
+        value = attn.head_to_batch_dim(value).contiguous()
+
+        hidden_states = xformers.ops.memory_efficient_attention(
+            query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale
+        )
+        hidden_states = hidden_states.to(query.dtype)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+class CrossAttnDownBlockSpatioTemporalWithFlow(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        temb_channels: int,
+        flow_channels: int,
+        num_layers: int = 1,
+        transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+        num_attention_heads: int = 1,
+        cross_attention_dim: int = 1280,
+        add_downsample: bool = True,
+        num_frames: int = 14,
+        pos_embed_dim: int = 64,
+        drag_token_cross_attn: bool = True,
+        use_modulate: bool = True,
+        drag_embedder_out_channels = (256, 320, 320),
+        num_max_drags: int = 5,
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+        flow_convs = []
+        if drag_token_cross_attn:
+            drag_token_mlps = []
+        self.num_max_drags = num_max_drags
+        self.num_frames = num_frames
+        self.pos_embed_dim = pos_embed_dim
+        self.drag_token_cross_attn = drag_token_cross_attn
+
+        self.has_cross_attention = True
+        self.num_attention_heads = num_attention_heads
+        self.use_modulate = use_modulate
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * num_layers
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                SpatioTemporalResBlock(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=1e-6,
+                )
+            )
+            attentions.append(
+                TransformerSpatioTemporalModel(
+                    num_attention_heads,
+                    out_channels // num_attention_heads,
+                    in_channels=out_channels,
+                    num_layers=transformer_layers_per_block[i],
+                    cross_attention_dim=cross_attention_dim,
+                )
+            )
+            flow_convs.append(
+                DragEmbedding(
+                    conditioning_channels=flow_channels, 
+                    conditioning_embedding_channels=out_channels * 2 if use_modulate else out_channels,
+                    block_out_channels = drag_embedder_out_channels,
+                )
+            )
+            if drag_token_cross_attn:
+                drag_token_mlps.append(
+                    nn.Sequential(
+                        nn.Linear(pos_embed_dim * 2 + out_channels * 2, cross_attention_dim),
+                        nn.SiLU(),
+                        nn.Linear(cross_attention_dim, cross_attention_dim),
+                    )
+                )
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+        self.flow_convs = nn.ModuleList(flow_convs)
+        if drag_token_cross_attn:
+            self.drag_token_mlps = nn.ModuleList(drag_token_mlps)
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample2D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        padding=1,
+                        name="op",
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.pos_embedding = {res: torch.tensor(get_2d_sincos_pos_embed(self.pos_embed_dim, res)) for res in [32, 16, 8, 4, 2]}
+        self.pos_embedding_prepared = False
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        temb: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        image_only_indicator: Optional[torch.Tensor] = None,
+        flow: Optional[torch.Tensor] = None,
+        drag_original: Optional[torch.Tensor] = None,  # (batch_frame, num_points, 4)
+    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
+        output_states = ()
+
+        batch_frame = hidden_states.shape[0]
+
+        if self.drag_token_cross_attn:
+            encoder_hidden_states_ori = encoder_hidden_states
+
+        if not self.pos_embedding_prepared:
+            for res in self.pos_embedding:
+                self.pos_embedding[res] = self.pos_embedding[res].to(hidden_states)
+            self.pos_embedding_prepared = True
+
+        blocks = list(zip(self.resnets, self.attentions, self.flow_convs))
+        for bid, (resnet, attn, flow_conv) in enumerate(blocks):
+            if self.training and self.gradient_checkpointing:  # TODO
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet),
+                    hidden_states,
+                    temb,
+                    image_only_indicator,
+                    **ckpt_kwargs,
+                )
+
+                if flow is not None:
+                    # flow shape is (batch_frame, 40, h, w)
+                    drags = flow.view(-1, self.num_frames, *flow.shape[1:])
+                    drags = drags.chunk(self.num_max_drags, dim=2)  # (batch, frame, 4, h, w) x 10
+                    drags = torch.stack(drags, dim=0)  # 10, batch, frame, 4, h, w
+                    invalid_flag = torch.all(drags == -1, dim=(2, 3, 4, 5))
+                    if self.use_modulate:
+                        scale, shift = flow_conv(flow).chunk(2, dim=1)
+                    else:
+                        scale = 0
+                        shift = flow_conv(flow)
+                    hidden_states = hidden_states * (1 + scale) + shift
+                    # print(self.drag_token_cross_attn)
+                    if self.drag_token_cross_attn:
+                        drag_token_mlp = self.drag_token_mlps[bid]
+                        pos_embed = self.pos_embedding[scale.shape[-1]]
+                        pos_embed = pos_embed.reshape(1, scale.shape[-1], scale.shape[-1], -1).permute(0, 3, 1, 2)
+                        grid = (drag_original[..., :2] * 2 - 1)[:, None]
+                        grid_end = (drag_original[..., 2:] * 2 - 1)[:, None]
+                        drags_pos_start = F.grid_sample(pos_embed.repeat(batch_frame, 1, 1, 1), grid, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        drags_pos_end = F.grid_sample(pos_embed.repeat(batch_frame, 1, 1, 1), grid_end, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        features = F.grid_sample(hidden_states.detach(), grid, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        features_end = F.grid_sample(hidden_states.detach(), grid_end, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+
+                        drag_token_in = torch.cat([features, features_end, drags_pos_start, drags_pos_end], dim=1).permute(0, 2, 1)
+                        drag_token_out = drag_token_mlp(drag_token_in)
+                        # Mask the invalid drags
+                        drag_token_out = drag_token_out.view(batch_frame // self.num_frames, self.num_frames, self.num_max_drags, -1)
+                        drag_token_out = drag_token_out.permute(2, 0, 1, 3)
+                        drag_token_out = drag_token_out.masked_fill(invalid_flag[..., None, None].expand_as(drag_token_out), 0)
+                        drag_token_out = drag_token_out.permute(1, 2, 0, 3).flatten(0, 1)
+                        encoder_hidden_states = torch.cat([encoder_hidden_states_ori, drag_token_out], dim=1)
+
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    image_only_indicator=image_only_indicator,
+                    return_dict=False,
+                )[0]
+            else:
+                hidden_states = resnet(
+                    hidden_states,
+                    temb,
+                    image_only_indicator=image_only_indicator,
+                )
+                if flow is not None:
+                    # flow shape is (batch_frame, 40, h, w)
+                    drags = flow.view(-1, self.num_frames, *flow.shape[1:])
+                    drags = drags.chunk(self.num_max_drags, dim=2)  # (batch, frame, 4, h, w) x 10
+                    drags = torch.stack(drags, dim=0)  # 10, batch, frame, 4, h, w
+                    invalid_flag = torch.all(drags == -1, dim=(2, 3, 4, 5))
+                    if self.use_modulate:
+                        scale, shift = flow_conv(flow).chunk(2, dim=1)
+                    else:
+                        scale = 0
+                        shift = flow_conv(flow)
+                    hidden_states = hidden_states * (1 + scale) + shift
+                    if self.drag_token_cross_attn:
+                        drag_token_mlp = self.drag_token_mlps[bid]
+                        pos_embed = self.pos_embedding[scale.shape[-1]]
+                        pos_embed = pos_embed.reshape(1, scale.shape[-1], scale.shape[-1], -1).permute(0, 3, 1, 2)
+                        grid = (drag_original[..., :2] * 2 - 1)[:, None]
+                        grid_end = (drag_original[..., 2:] * 2 - 1)[:, None]
+                        drags_pos_start = F.grid_sample(pos_embed.repeat(batch_frame, 1, 1, 1), grid, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        drags_pos_end = F.grid_sample(pos_embed.repeat(batch_frame, 1, 1, 1), grid_end, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        features = F.grid_sample(hidden_states.detach(), grid, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        features_end = F.grid_sample(hidden_states.detach(), grid_end, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+
+                        drag_token_in = torch.cat([features, features_end, drags_pos_start, drags_pos_end], dim=1).permute(0, 2, 1)
+                        drag_token_out = drag_token_mlp(drag_token_in)
+                        # Mask the invalid drags
+                        drag_token_out = drag_token_out.view(batch_frame // self.num_frames, self.num_frames, self.num_max_drags, -1)
+                        drag_token_out = drag_token_out.permute(2, 0, 1, 3)
+                        drag_token_out = drag_token_out.masked_fill(invalid_flag[..., None, None].expand_as(drag_token_out), 0)
+                        drag_token_out = drag_token_out.permute(1, 2, 0, 3).flatten(0, 1)
+                        encoder_hidden_states = torch.cat([encoder_hidden_states_ori, drag_token_out], dim=1)
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    image_only_indicator=image_only_indicator,
+                    return_dict=False,
+                )[0]
+
+            output_states = output_states + (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states = output_states + (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class CrossAttnUpBlockSpatioTemporalWithFlow(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        prev_output_channel: int,
+        temb_channels: int,
+        flow_channels: int,
+        resolution_idx: Optional[int] = None,
+        num_layers: int = 1,
+        transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+        resnet_eps: float = 1e-6,
+        num_attention_heads: int = 1,
+        cross_attention_dim: int = 1280,
+        add_upsample: bool = True,
+        num_frames: int = 14,
+        pos_embed_dim: int = 64,
+        drag_token_cross_attn: bool = True,
+        use_modulate: bool = True,
+        drag_embedder_out_channels = (256, 320, 320),
+        num_max_drags: int = 5,
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+        flow_convs = []
+        if drag_token_cross_attn:
+            drag_token_mlps = []
+        self.num_max_drags = num_max_drags
+
+        self.drag_token_cross_attn = drag_token_cross_attn
+
+        self.num_frames = num_frames
+        self.pos_embed_dim = pos_embed_dim
+
+        self.has_cross_attention = True
+        self.num_attention_heads = num_attention_heads
+        self.use_modulate = use_modulate
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * num_layers
+
+        for i in range(num_layers):
+            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
+            resnet_in_channels = prev_output_channel if i == 0 else out_channels
+
+            resnets.append(
+                SpatioTemporalResBlock(
+                    in_channels=resnet_in_channels + res_skip_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                )
+            )
+            attentions.append(
+                TransformerSpatioTemporalModel(
+                    num_attention_heads,
+                    out_channels // num_attention_heads,
+                    in_channels=out_channels,
+                    num_layers=transformer_layers_per_block[i],
+                    cross_attention_dim=cross_attention_dim,
+                )
+            )
+            flow_convs.append(
+                DragEmbedding(
+                    conditioning_channels=flow_channels, 
+                    conditioning_embedding_channels=out_channels * 2 if use_modulate else out_channels,
+                    block_out_channels = drag_embedder_out_channels,
+                )
+            )
+            if drag_token_cross_attn:
+                drag_token_mlps.append(
+                    nn.Sequential(
+                        nn.Linear(pos_embed_dim * 2 + out_channels * 2, cross_attention_dim),
+                        nn.SiLU(),
+                        nn.Linear(cross_attention_dim, cross_attention_dim),
+                    )
+                )
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+        self.flow_convs = nn.ModuleList(flow_convs)
+        
+        if drag_token_cross_attn:
+            self.drag_token_mlps = nn.ModuleList(drag_token_mlps)
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+        self.pos_embedding = {res: torch.tensor(get_2d_sincos_pos_embed(pos_embed_dim, res)) for res in [32, 16, 8, 4, 2]}
+        self.pos_embedding_prepared = False
+
+        self.gradient_checkpointing = False
+        self.resolution_idx = resolution_idx
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
+        temb: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        image_only_indicator: Optional[torch.Tensor] = None,
+        flow: Optional[torch.Tensor] = None,
+        drag_original: Optional[torch.Tensor] = None,  # (batch_frame, num_points, 4)
+    ) -> torch.FloatTensor:
+        batch_frame = hidden_states.shape[0]
+
+        if self.drag_token_cross_attn:
+            encoder_hidden_states_ori = encoder_hidden_states
+        
+        if not self.pos_embedding_prepared:
+            for res in self.pos_embedding:
+                self.pos_embedding[res] = self.pos_embedding[res].to(hidden_states)
+            self.pos_embedding_prepared = True
+
+        for bid, (resnet, attn, flow_conv) in enumerate(zip(self.resnets, self.attentions, self.flow_convs)):
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.training and self.gradient_checkpointing:  # TODO
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet),
+                    hidden_states,
+                    temb,
+                    image_only_indicator,
+                    **ckpt_kwargs,
+                )
+                if flow is not None:
+                    # flow shape is (batch_frame, 40, h, w)
+                    drags = flow.view(-1, self.num_frames, *flow.shape[1:])
+                    drags = drags.chunk(self.num_max_drags, dim=2)  # (batch, frame, 4, h, w) x 10
+                    drags = torch.stack(drags, dim=0)  # 10, batch, frame, 4, h, w
+                    invalid_flag = torch.all(drags == -1, dim=(2, 3, 4, 5))
+                    if self.use_modulate:
+                        scale, shift = flow_conv(flow).chunk(2, dim=1)
+                    else:
+                        scale = 0
+                        shift = flow_conv(flow)
+                    hidden_states = hidden_states * (1 + scale) + shift
+                    if self.drag_token_cross_attn:
+                        drag_token_mlp = self.drag_token_mlps[bid]
+                        pos_embed = self.pos_embedding[scale.shape[-1]]
+                        pos_embed = pos_embed.reshape(1, scale.shape[-1], scale.shape[-1], -1).permute(0, 3, 1, 2)
+                        grid = (drag_original[..., :2] * 2 - 1)[:, None]
+                        grid_end = (drag_original[..., 2:] * 2 - 1)[:, None]
+                        drags_pos_start = F.grid_sample(pos_embed.repeat(batch_frame, 1, 1, 1), grid, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        drags_pos_end = F.grid_sample(pos_embed.repeat(batch_frame, 1, 1, 1), grid_end, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        features = F.grid_sample(hidden_states.detach(), grid, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        features_end = F.grid_sample(hidden_states.detach(), grid_end, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+
+                        drag_token_in = torch.cat([features, features_end, drags_pos_start, drags_pos_end], dim=1).permute(0, 2, 1)
+                        drag_token_out = drag_token_mlp(drag_token_in)
+                        # Mask the invalid drags
+                        drag_token_out = drag_token_out.view(batch_frame // self.num_frames, self.num_frames, self.num_max_drags, -1)
+                        drag_token_out = drag_token_out.permute(2, 0, 1, 3)
+                        drag_token_out = drag_token_out.masked_fill(invalid_flag[..., None, None].expand_as(drag_token_out), 0)
+                        drag_token_out = drag_token_out.permute(1, 2, 0, 3).flatten(0, 1)
+                        encoder_hidden_states = torch.cat([encoder_hidden_states_ori, drag_token_out], dim=1)
+
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    image_only_indicator=image_only_indicator,
+                    return_dict=False,
+                )[0]
+            else:
+                hidden_states = resnet(
+                    hidden_states,
+                    temb,
+                    image_only_indicator=image_only_indicator,
+                )
+                if flow is not None:
+                    # flow shape is (batch_frame, 40, h, w)
+                    drags = flow.view(-1, self.num_frames, *flow.shape[1:])
+                    drags = drags.chunk(self.num_max_drags, dim=2)  # (batch, frame, 4, h, w) x 10
+                    drags = torch.stack(drags, dim=0)  # 10, batch, frame, 4, h, w
+                    invalid_flag = torch.all(drags == -1, dim=(2, 3, 4, 5))
+                    if self.use_modulate:
+                        scale, shift = flow_conv(flow).chunk(2, dim=1)
+                    else:
+                        scale = 0
+                        shift = flow_conv(flow)
+                    hidden_states = hidden_states * (1 + scale) + shift
+                    if self.drag_token_cross_attn:
+                        drag_token_mlp = self.drag_token_mlps[bid]
+                        pos_embed = self.pos_embedding[scale.shape[-1]]
+                        pos_embed = pos_embed.reshape(1, scale.shape[-1], scale.shape[-1], -1).permute(0, 3, 1, 2)
+                        grid = (drag_original[..., :2] * 2 - 1)[:, None]
+                        grid_end = (drag_original[..., 2:] * 2 - 1)[:, None]
+                        drags_pos_start = F.grid_sample(pos_embed.repeat(batch_frame, 1, 1, 1), grid, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        drags_pos_end = F.grid_sample(pos_embed.repeat(batch_frame, 1, 1, 1), grid_end, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        features = F.grid_sample(hidden_states.detach(), grid, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+                        features_end = F.grid_sample(hidden_states.detach(), grid_end, padding_mode="border", mode="bilinear", align_corners=False).squeeze(dim=2)
+
+                        drag_token_in = torch.cat([features, features_end, drags_pos_start, drags_pos_end], dim=1).permute(0, 2, 1)
+                        drag_token_out = drag_token_mlp(drag_token_in)
+                        # Mask the invalid drags
+                        drag_token_out = drag_token_out.view(batch_frame // self.num_frames, self.num_frames, self.num_max_drags, -1)
+                        drag_token_out = drag_token_out.permute(2, 0, 1, 3)
+                        drag_token_out = drag_token_out.masked_fill(invalid_flag[..., None, None].expand_as(drag_token_out), 0)
+                        drag_token_out = drag_token_out.permute(1, 2, 0, 3).flatten(0, 1)
+                        encoder_hidden_states = torch.cat([encoder_hidden_states_ori, drag_token_out], dim=1)
+
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    image_only_indicator=image_only_indicator,
+                    return_dict=False,
+                )[0]
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        return hidden_states
+
+
+def get_down_block(
+    with_concatenated_flow: bool = False,
+    *args,
+    **kwargs,
+):
+    NEEDED_KEYS = [
+        "in_channels",
+        "out_channels",
+        "temb_channels",
+        "flow_channels",
+        "num_layers",
+        "transformer_layers_per_block",
+        "num_attention_heads",
+        "cross_attention_dim",
+        "add_downsample",
+        "pos_embed_dim",
+        'use_modulate',
+        "drag_token_cross_attn",
+        "drag_embedder_out_channels",
+        "num_max_drags",
+    ]
+    if not with_concatenated_flow or args[0] == "DownBlockSpatioTemporal":
+        kwargs.pop("flow_channels", None)
+        kwargs.pop("pos_embed_dim", None)
+        kwargs.pop("use_modulate", None)
+        kwargs.pop("drag_token_cross_attn", None)
+        kwargs.pop("drag_embedder_out_channels", None)
+        kwargs.pop("num_max_drags", None)
+        return gdb(*args, **kwargs)
+    elif args[0] == "CrossAttnDownBlockSpatioTemporal":
+        for key in list(kwargs.keys()):
+            if key not in NEEDED_KEYS:
+                kwargs.pop(key, None)
+        return CrossAttnDownBlockSpatioTemporalWithFlow(*args[1:], **kwargs)
+    else:
+        raise ValueError(f"Unknown block type {args[0]}")
+    
+
+def get_up_block(
+    with_concatenated_flow: bool = False,
+    *args,
+    **kwargs,
+):
+    NEEDED_KEYS = [
+        "in_channels",
+        "out_channels",
+        "prev_output_channel",
+        "temb_channels",
+        "flow_channels",
+        "resolution_idx",
+        "num_layers",
+        "transformer_layers_per_block",
+        "resnet_eps",
+        "num_attention_heads",
+        "cross_attention_dim",
+        "add_upsample",
+        "pos_embed_dim",
+        "use_modulate",
+        "drag_token_cross_attn",
+        "drag_embedder_out_channels",
+        "num_max_drags",
+    ]
+    if not with_concatenated_flow or args[0] == "UpBlockSpatioTemporal":
+        kwargs.pop("flow_channels", None)
+        kwargs.pop("pos_embed_dim", None)
+        kwargs.pop("use_modulate", None)
+        kwargs.pop("drag_token_cross_attn", None)
+        kwargs.pop("drag_embedder_out_channels", None)
+        kwargs.pop("num_max_drags", None)
+        return gub(*args, **kwargs)
+    elif args[0] == "CrossAttnUpBlockSpatioTemporal":
+        for key in list(kwargs.keys()):
+            if key not in NEEDED_KEYS:
+                kwargs.pop(key, None)
+        return CrossAttnUpBlockSpatioTemporalWithFlow(*args[1:], **kwargs)
+    else:
+        raise ValueError(f"Unknown block type {args[0]}")
+
+
+@dataclass
+class UNetSpatioTemporalConditionOutput(BaseOutput):
+    """
+    The output of [`UNetSpatioTemporalConditionModel`].
+
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`):
+            The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
+    """
+
+    sample: torch.FloatTensor = None
+
+
+class UNetDragSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
+    r"""
+    A conditional Spatio-Temporal UNet model that takes a noisy video frames, conditional state, and a timestep and
+    returns a sample shaped output.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+            Height and width of input/output sample.
+        in_channels (`int`, *optional*, defaults to 8): Number of channels in the input sample.
+        out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`):
+            The tuple of downsample blocks to use.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`):
+            The tuple of upsample blocks to use.
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+            The tuple of output channels for each block.
+        addition_time_embed_dim: (`int`, defaults to 256):
+            Dimension to to encode the additional time ids.
+        projection_class_embeddings_input_dim (`int`, defaults to 768):
+            The dimension of the projection of encoded `added_time_ids`.
+        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
+        cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+            The dimension of the cross attention features.
+        transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
+            [`~models.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal`],
+            [`~models.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal`],
+            [`~models.unet_3d_blocks.UNetMidBlockSpatioTemporal`].
+        num_attention_heads (`int`, `Tuple[int]`, defaults to `(5, 10, 10, 20)`):
+            The number of attention heads.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 8,
+        out_channels: int = 4,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlockSpatioTemporal",
+            "CrossAttnDownBlockSpatioTemporal",
+            "CrossAttnDownBlockSpatioTemporal",
+            "DownBlockSpatioTemporal",
+        ),
+        up_block_types: Tuple[str] = (
+            "UpBlockSpatioTemporal",
+            "CrossAttnUpBlockSpatioTemporal",
+            "CrossAttnUpBlockSpatioTemporal",
+            "CrossAttnUpBlockSpatioTemporal",
+        ),
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        addition_time_embed_dim: int = 256,
+        projection_class_embeddings_input_dim: int = 768,
+        layers_per_block: Union[int, Tuple[int]] = 2,
+        cross_attention_dim: Union[int, Tuple[int]] = 1024,
+        transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
+        num_attention_heads: Union[int, Tuple[int]] = (5, 10, 20, 20),
+        num_frames: int = 25,
+        num_drags: int = 10,
+        cond_dropout_prob: float = 0.1,
+        pos_embed_dim: int = 64,
+        drag_token_cross_attn: bool = True,
+
+        use_modulate: bool = True,
+
+        drag_embedder_out_channels = (256, 320, 320),
+
+        cross_attn_with_ref: bool = True,
+        double_batch: bool = False,
+    ):
+        super().__init__()
+
+        self.sample_size = sample_size
+        self.cond_dropout_prob = cond_dropout_prob
+        self.drag_token_cross_attn = drag_token_cross_attn
+
+        self.pos_embed_dim = pos_embed_dim
+
+        self.use_modulate = use_modulate
+
+        self.cross_attn_with_ref = cross_attn_with_ref
+        self.double_batch = double_batch
+
+        flow_channels = 6 * num_drags
+
+        # Check inputs
+        if len(down_block_types) != len(up_block_types):
+            raise ValueError(
+                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+            )
+
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
+            )
+
+        if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
+            )
+
+        # input
+        self.conv_in = nn.Conv2d(
+            in_channels,
+            block_out_channels[0],
+            kernel_size=3,
+            padding=1,
+        )
+
+        # time
+        time_embed_dim = block_out_channels[0] * 4
+
+        self.time_proj = Timesteps(block_out_channels[0], True, downscale_freq_shift=0)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+
+        self.down_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        if isinstance(num_attention_heads, int):
+            num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+        if isinstance(cross_attention_dim, int):
+            cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
+
+        if isinstance(layers_per_block, int):
+            layers_per_block = [layers_per_block] * len(down_block_types)
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+        blocks_time_embed_dim = time_embed_dim
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                True,
+                down_block_type,
+                num_layers=layers_per_block[i],
+                transformer_layers_per_block=transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=1e-5,
+                cross_attention_dim=cross_attention_dim[i],
+                num_attention_heads=num_attention_heads[i],
+                resnet_act_fn="silu",
+                flow_channels=flow_channels,
+                pos_embed_dim=pos_embed_dim,
+                use_modulate=use_modulate,
+                drag_token_cross_attn=drag_token_cross_attn,
+                drag_embedder_out_channels=drag_embedder_out_channels,
+                num_max_drags=num_drags,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = UNetMidBlockSpatioTemporal(
+            block_out_channels[-1],
+            temb_channels=blocks_time_embed_dim,
+            transformer_layers_per_block=transformer_layers_per_block[-1],
+            cross_attention_dim=cross_attention_dim[-1],
+            num_attention_heads=num_attention_heads[-1],
+        )
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_num_attention_heads = list(reversed(num_attention_heads))
+        reversed_layers_per_block = list(reversed(layers_per_block))
+        reversed_cross_attention_dim = list(reversed(cross_attention_dim))
+        reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
+
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            is_final_block = i == len(block_out_channels) - 1
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            # add upsample block for all BUT final layer
+            if not is_final_block:
+                add_upsample = True
+                self.num_upsamplers += 1
+            else:
+                add_upsample = False
+
+            up_block = get_up_block(
+                True,
+                up_block_type,
+                num_layers=reversed_layers_per_block[i] + 1,
+                transformer_layers_per_block=reversed_transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=1e-5,
+                resolution_idx=i,
+                cross_attention_dim=reversed_cross_attention_dim[i],
+                num_attention_heads=reversed_num_attention_heads[i],
+                resnet_act_fn="silu",
+                flow_channels=flow_channels,
+                pos_embed_dim=pos_embed_dim,
+                use_modulate=use_modulate,
+                drag_token_cross_attn=drag_token_cross_attn,
+                drag_embedder_out_channels=drag_embedder_out_channels,
+                num_max_drags=num_drags,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-5)
+        self.conv_act = nn.SiLU()
+
+        self.conv_out = nn.Conv2d(
+            block_out_channels[0],
+            out_channels,
+            kernel_size=3,
+            padding=1,
+        )
+
+        self.num_drags = num_drags
+
+        self.pos_embedding = {res: torch.tensor(get_2d_sincos_pos_embed(self.pos_embed_dim, res)) for res in [32, 16, 8, 4, 2]}
+        self.pos_embedding_prepared = False
+
+    @property
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+            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: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        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 set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
+    def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+        """
+        Sets the attention processor to use [feed forward
+        chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
+
+        Parameters:
+            chunk_size (`int`, *optional*):
+                The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
+                over each tensor of dim=`dim`.
+            dim (`int`, *optional*, defaults to `0`):
+                The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
+                or dim=1 (sequence length).
+        """
+        if dim not in [0, 1]:
+            raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
+
+        # By default chunk size is 1
+        chunk_size = chunk_size or 1
+
+        def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
+            if hasattr(module, "set_chunk_feed_forward"):
+                module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
+
+            for child in module.children():
+                fn_recursive_feed_forward(child, chunk_size, dim)
+
+        for module in self.children():
+            fn_recursive_feed_forward(module, chunk_size, dim)
+
+    def _convert_drag_to_concatting_image(self, drags: torch.Tensor, current_resolution: int) -> torch.Tensor:
+        batch_size, num_frames, num_points, _ = drags.shape
+        num_channels = 6
+        concatting_image = -torch.ones(
+            batch_size, num_frames, num_channels * num_points, current_resolution, current_resolution
+        ).to(drags)
+        
+        not_all_zeros = drags.any(dim=-1).repeat_interleave(num_channels, dim=-1)[..., None, None]
+        y_grid, x_grid = torch.meshgrid(torch.arange(current_resolution), torch.arange(current_resolution), indexing='ij')
+        y_grid = y_grid.to(drags)[None, None, None]  # (1, 1, 1, res, res)
+        x_grid = x_grid.to(drags)[None, None, None]  # (1, 1, 1, res, res)
+        x0 = (drags[..., 0] * current_resolution - 0.5).round().clip(0, current_resolution - 1)
+        x_src = (drags[..., 0] * current_resolution - x0)[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        x0 = x0[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        x0 = torch.stack([
+            x0, x0, 
+            torch.zeros_like(x0) - 1, torch.zeros_like(x0) - 1, 
+            torch.zeros_like(x0) - 1, torch.zeros_like(x0) - 1,
+        ], dim=3).view(batch_size, num_frames, num_channels * num_points, 1, 1)
+
+        y0 = (drags[..., 1] * current_resolution - 0.5).round().clip(0, current_resolution - 1)
+        y_src = (drags[..., 1] * current_resolution - y0)[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        y0 = y0[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        y0 = torch.stack([
+            y0, y0, 
+            torch.zeros_like(y0) - 1, torch.zeros_like(y0) - 1, 
+            torch.zeros_like(y0) - 1, torch.zeros_like(y0) - 1,
+        ], dim=3).view(batch_size, num_frames, num_channels * num_points, 1, 1)
+
+        x1 = (drags[..., 2] * current_resolution - 0.5).round().clip(0, current_resolution - 1)
+        x_tgt = (drags[..., 2] * current_resolution - x1)[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        x1 = x1[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        x1 = torch.stack([
+            torch.zeros_like(x1) - 1, torch.zeros_like(x1) - 1, 
+            x1, x1, 
+            torch.zeros_like(x1) - 1, torch.zeros_like(x1) - 1
+        ], dim=3).view(batch_size, num_frames, num_channels * num_points, 1, 1)
+
+        y1 = (drags[..., 3] * current_resolution - 0.5).round().clip(0, current_resolution - 1)
+        y_tgt = (drags[..., 3] * current_resolution - y1)[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        y1 = y1[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        y1 = torch.stack([
+            torch.zeros_like(y1) - 1, torch.zeros_like(y1) - 1, 
+            y1, y1, 
+            torch.zeros_like(y1) - 1, torch.zeros_like(y1) - 1
+        ], dim=3).view(batch_size, num_frames, num_channels * num_points, 1, 1)
+
+        drags_final = drags[:, -1:, :, :].expand_as(drags)
+        x_final = (drags_final[..., 2] * current_resolution - 0.5).round().clip(0, current_resolution - 1)
+        x_final_tgt = (drags_final[..., 2] * current_resolution - x_final)[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        x_final = x_final[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        x_final = torch.stack([
+            torch.zeros_like(x_final) - 1, torch.zeros_like(x_final) - 1, 
+            torch.zeros_like(x_final) - 1, torch.zeros_like(x_final) - 1, 
+            x_final, x_final
+        ], dim=3).view(batch_size, num_frames, num_channels * num_points, 1, 1)
+
+        y_final = (drags_final[..., 3] * current_resolution - 0.5).round().clip(0, current_resolution - 1)
+        y_final_tgt = (drags_final[..., 3] * current_resolution - y_final)[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        y_final = y_final[..., None, None]  # (batch, num_frames, num_points, 1, 1)
+        y_final = torch.stack([
+            torch.zeros_like(y_final) - 1, torch.zeros_like(y_final) - 1, 
+            torch.zeros_like(y_final) - 1, torch.zeros_like(y_final) - 1, 
+            y_final, y_final
+        ], dim=3).view(batch_size, num_frames, num_channels * num_points, 1, 1)
+
+        value_image = torch.stack([
+            x_src, y_src, 
+            x_tgt, y_tgt, 
+            x_final_tgt, y_final_tgt
+        ], dim=3).view(batch_size, num_frames, num_channels * num_points, 1, 1)
+        value_image = value_image.expand_as(concatting_image)
+        start_mask = (x_grid == x0) & (y_grid == y0) & not_all_zeros
+        end_mask = (x_grid == x1) & (y_grid == y1) & not_all_zeros
+        final_mask = (x_grid == x_final) & (y_grid == y_final) & not_all_zeros
+        concatting_image[start_mask] = value_image[start_mask]
+        concatting_image[end_mask] = value_image[end_mask]
+        concatting_image[final_mask] = value_image[final_mask]
+        return concatting_image
+    
+    def zero_init(self):
+        for block in self.down_blocks:
+            if hasattr(block, "flow_convs"):
+                for flow_conv in block.flow_convs:
+                    try:
+                        nn.init.constant_(flow_conv.conv_out.weight, 0)
+                        nn.init.constant_(flow_conv.conv_out.bias, 0)
+                    except:
+                        nn.init.constant_(flow_conv.weight, 0)
+
+        for block in self.up_blocks:
+            if hasattr(block, "flow_convs"):
+                for flow_conv in block.flow_convs:
+                    try:
+                        nn.init.constant_(flow_conv.conv_out.weight, 0)
+                        nn.init.constant_(flow_conv.conv_out.bias, 0)
+                    except:
+                        nn.init.constant_(flow_conv.weight, 0)
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        image_latents: torch.FloatTensor,
+        encoder_hidden_states: torch.Tensor,
+        added_time_ids: torch.Tensor,
+        drags: torch.Tensor,
+
+        force_drop_ids: Optional[torch.Tensor] = None,
+    ) -> Union[UNetSpatioTemporalConditionOutput, Tuple]:
+        r"""
+        The [`UNetSpatioTemporalConditionModel`] forward method.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The noisy input tensor with the following shape `(batch, num_frames, channel, height, width)`.
+            image_latents (`torch.FloatTensor`):
+                The clean conditioning tensor of the first frame of the image with shape `(batch, num_channels, height, width)`.
+            timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
+            encoder_hidden_states (`torch.FloatTensor`):
+                The encoder hidden states with shape `(batch, sequence_length, cross_attention_dim)`.
+            added_time_ids: (`torch.FloatTensor`):
+                The additional time ids with shape `(batch, num_additional_ids)`. These are encoded with sinusoidal
+                embeddings and added to the time embeddings.
+            drags (`torch.Tensor`):
+                The drags tensor with shape `(batch, num_frames, num_points, 4)`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] instead
+                of a plain tuple.
+        Returns:
+            [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] or `tuple`:
+                If `return_dict` is True, an [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] is
+                returned, otherwise a `tuple` is returned where the first element is the sample tensor.
+        """
+        batch_size, num_frames = sample.shape[:2]
+
+        if not self.pos_embedding_prepared:
+            for res in self.pos_embedding:
+                self.pos_embedding[res] = self.pos_embedding[res].to(drags)
+            self.pos_embedding_prepared = True
+
+        # 0. prepare for cfg
+        drag_drop_ids = None
+        if (self.training and self.cond_dropout_prob > 0) or force_drop_ids is not None:
+            if force_drop_ids is None:
+                drag_drop_ids = torch.rand(batch_size, device=sample.device) < self.cond_dropout_prob
+            else:
+                drag_drop_ids = (force_drop_ids == 1)
+            drags = drags * ~drag_drop_ids[:, None, None, None]
+
+        sample = torch.cat([sample, image_latents[:, None].repeat(1, num_frames, 1, 1, 1)], dim=2)
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(batch_size)
+
+        if self.cross_attn_with_ref and self.double_batch:
+            sample_ref = image_latents[:, None].repeat(1, num_frames, 2, 1, 1)
+            sample_ref[:, :, :4] = sample_ref[:, :, :4] * 0.18215
+            sample = torch.cat([sample_ref, sample], dim=0)
+
+            drags = torch.cat([torch.zeros_like(drags), drags], dim=0)
+            encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states], dim=0)
+            timesteps = torch.cat([timesteps, timesteps], dim=0)
+            batch_size *= 2
+        
+        drag_encodings = {res: self._convert_drag_to_concatting_image(drags, res) for res in [32, 16, 8]}
+
+        t_emb = self.time_proj(timesteps)
+
+        # `Timesteps` does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=sample.dtype)
+        emb = self.time_embedding(t_emb)
+
+        # Flatten the batch and frames dimensions
+        # sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width]
+        sample = sample.flatten(0, 1)
+        # Repeat the embeddings num_video_frames times
+        # emb: [batch, channels] -> [batch * frames, channels]
+        emb = emb.repeat_interleave(num_frames, dim=0)
+        # encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels]
+        encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0)
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device)
+
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                flow = drag_encodings[sample.shape[-1]]
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    image_only_indicator=image_only_indicator,
+                    flow=flow.flatten(0, 1),
+                    drag_original=drags.flatten(0, 1),
+                )
+            else:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    image_only_indicator=image_only_indicator,
+                )
+
+            down_block_res_samples += res_samples
+
+        # 4. mid
+        sample = self.mid_block(
+            hidden_states=sample,
+            temb=emb,
+            encoder_hidden_states=encoder_hidden_states,
+            image_only_indicator=image_only_indicator,
+        )
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                flow = drag_encodings[sample.shape[-1]]
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    image_only_indicator=image_only_indicator,
+                    flow=flow.flatten(0, 1),
+                    drag_original=drags.flatten(0, 1),
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    image_only_indicator=image_only_indicator,
+                )
+
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        # 7. Reshape back to original shape
+        sample = sample.reshape(batch_size, num_frames, *sample.shape[1:])
+        if self.cross_attn_with_ref and self.double_batch:
+            sample = sample[batch_size // 2:]
+
+        return sample
+
+
+if __name__ == "__main__":
+    puppet_master = UNetDragSpatioTemporalConditionModel(num_drags=5)
+    state_dict = torch.load("ckpts/0800000-ema.pt", map_location="cpu")
+    puppet_master.load_state_dict(state_dict, strict=True)