[ { "path": "README.md", "content": "# MotionClone\nThis repository is the official implementation of [MotionClone](https://arxiv.org/abs/2406.05338). It is a **training-free framework** that enables motion cloning from a reference video for controllable video generation, **without cumbersome video inversion processes**.\n
Click for the full abstract of MotionClone\n\n> Motion-based controllable video generation offers the potential for creating captivating visual content. Existing methods typically necessitate model training to encode particular motion cues or incorporate fine-tuning to inject certain motion patterns, resulting in limited flexibility and generalization.\nIn this work, we propose **MotionClone** a training-free framework that enables motion cloning from reference videos to versatile motion-controlled video generation, including text-to-video and image-to-video. Based on the observation that the dominant components in temporal-attention maps drive motion synthesis, while the rest mainly capture noisy or very subtle motions, MotionClone utilizes sparse temporal attention weights as motion representations for motion guidance, facilitating diverse motion transfer across varying scenarios. Meanwhile, MotionClone allows for the direct extraction of motion representation through a single denoising step, bypassing the cumbersome inversion processes and thus promoting both efficiency and flexibility. \nExtensive experiments demonstrate that MotionClone exhibits proficiency in both global camera motion and local object motion, with notable superiority in terms of motion fidelity, textual alignment, and temporal consistency.\n
\n\n**[MotionClone: Training-Free Motion Cloning for Controllable Video Generation](https://arxiv.org/abs/2406.05338)** \n
\n[Pengyang Ling*](https://github.com/LPengYang/),\n[Jiazi Bu*](https://github.com/Bujiazi/),\n[Pan Zhang†](https://panzhang0212.github.io/),\n[Xiaoyi Dong](https://scholar.google.com/citations?user=FscToE0AAAAJ&hl=en/),\n[Yuhang Zang](https://yuhangzang.github.io/),\n[Tong Wu](https://wutong16.github.io/),\n[Huaian Chen](https://scholar.google.com.hk/citations?hl=zh-CN&user=D6ol9XkAAAAJ),\n[Jiaqi Wang](https://myownskyw7.github.io/),\n[Yi Jin†](https://scholar.google.ca/citations?hl=en&user=mAJ1dCYAAAAJ) \n(*Equal Contribution)(†Corresponding Author)\n\n\n[![arXiv](https://img.shields.io/badge/arXiv-2406.05338-b31b1b.svg)](https://arxiv.org/abs/2406.05338)\n[![Project Page](https://img.shields.io/badge/Project-Website-green)](https://bujiazi.github.io/motionclone.github.io/)\n![](https://img.shields.io/github/stars/LPengYang/MotionClone?style=social)\n\n\n\n## Demo\n[![]](https://github.com/user-attachments/assets/d1f1c753-f192-455b-9779-94c925e51aaa)\n\n\n## 🖋 News\n- The latest version of our paper (**v4**) is available on arXiv! (10.08)\n- The latest version of our paper (**v3**) is available on arXiv! (7.2)\n- Code released! (6.29)\n\n## 🏗️ Todo\n- [x] We have updated the latest version of MotionCloning, which performs motion transfer **without video inversion** and supports **image-to-video and sketch-to-video**.\n- [x] Release the MotionClone code (We have released **the first version** of our code and will continue to optimize it. We welcome any questions or issues you may have and will address them promptly.)\n- [x] Release paper\n\n## 📚 Gallery\nWe show more results in the [Project Page](https://bujiazi.github.io/motionclone.github.io/).\n\n## 🚀 Method Overview\n### Feature visualization\n
\n \n
\n\n### Pipeline\n
\n \n
\n\nMotionClone utilizes sparse temporal attention weights as motion representations for motion guidance, facilitating diverse motion transfer across varying scenarios. Meanwhile, MotionClone allows for the direct extraction of motion representation through a single denoising step, bypassing the cumbersome inversion processes and thus promoting both efficiency and flexibility.\n\n## 🔧 Installations (python==3.11.3 recommended)\n\n### Setup repository and conda environment\n\n```\ngit clone https://github.com/Bujiazi/MotionClone.git\ncd MotionClone\n\nconda env create -f environment.yaml\nconda activate motionclone\n```\n\n## 🔑 Pretrained Model Preparations\n\n### Download Stable Diffusion V1.5\n\n```\ngit lfs install\ngit clone https://huggingface.co/runwayml/stable-diffusion-v1-5 models/StableDiffusion/\n```\n\nAfter downloading Stable Diffusion, save them to `models/StableDiffusion`. \n\n### Prepare Community Models\n\nManually download the community `.safetensors` models from [RealisticVision V5.1](https://civitai.com/models/4201?modelVersionId=130072) and save them to `models/DreamBooth_LoRA`. \n\n### Prepare AnimateDiff Motion Modules\n\nManually download the AnimateDiff modules from [AnimateDiff](https://github.com/guoyww/AnimateDiff), we recommend [`v3_adapter_sd_v15.ckpt`](https://huggingface.co/guoyww/animatediff/blob/main/v3_sd15_adapter.ckpt) and [`v3_sd15_mm.ckpt.ckpt`](https://huggingface.co/guoyww/animatediff/blob/main/v3_sd15_mm.ckpt). Save the modules to `models/Motion_Module`.\n\n### Prepare SparseCtrl for image-to-video and sketch-to-video\nManually download \"v3_sd15_sparsectrl_rgb.ckpt\" and \"v3_sd15_sparsectrl_scribble.ckpt\" from [AnimateDiff](https://huggingface.co/guoyww/animatediff/tree/main). Save the modules to `models/SparseCtrl`.\n\n## 🎈 Quick Start\n\n### Perform Text-to-video generation with customized camera motion\n```\npython t2v_video_sample.py --inference_config \"configs/t2v_camera.yaml\" --examples \"configs/t2v_camera.jsonl\"\n```\n### Perform Text-to-video generation with customized object motion\n```\npython t2v_video_sample.py --inference_config \"configs/t2v_object.yaml\" --examples \"configs/t2v_object.jsonl\"\n```\n### Combine motion cloning with sketch-to-video\n```\npython i2v_video_sample.py --inference_config \"configs/i2v_sketch.yaml\" --examples \"configs/i2v_sketch.jsonl\"\n```\n### Combine motion cloning with image-to-video\n```\npython i2v_video_sample.py --inference_config \"configs/i2v_rgb.yaml\" --examples \"configs/i2v_rgb.jsonl\"\n```\n\n\n## 📎 Citation \n\nIf you find this work helpful, please cite the following paper:\n\n```\n@article{ling2024motionclone,\n title={MotionClone: Training-Free Motion Cloning for Controllable Video Generation},\n author={Ling, Pengyang and Bu, Jiazi and Zhang, Pan and Dong, Xiaoyi and Zang, Yuhang and Wu, Tong and Chen, Huaian and Wang, Jiaqi and Jin, Yi},\n journal={arXiv preprint arXiv:2406.05338},\n year={2024}\n}\n```\n\n## 📣 Disclaimer\n\nThis is official code of MotionClone.\nAll the copyrights of the demo images and audio are from community users. \nFeel free to contact us if you would like remove them.\n\n## 💞 Acknowledgements\nThe code is built upon the below repositories, we thank all the contributors for open-sourcing.\n* [AnimateDiff](https://github.com/guoyww/AnimateDiff)\n* [FreeControl](https://github.com/genforce/freecontrol)\n" }, { "path": "configs/i2v_rgb.jsonl", "content": "{\"video_path\":\"reference_videos/camera_zoom_out.mp4\", \"condition_image_paths\":[\"condition_images/rgb/dog_on_grass.png\"], \"new_prompt\": \"Dog, lying on the grass\"}" }, { "path": "configs/i2v_rgb.yaml", "content": "motion_module: \"models/Motion_Module/v3_sd15_mm.ckpt\"\ndreambooth_path: \"models/DreamBooth_LoRA/realisticVisionV60B1_v51VAE.safetensors\"\nmodel_config: \"configs/model_config/model_config.yaml\"\ncontrolnet_path: \"models/SparseCtrl/v3_sd15_sparsectrl_rgb.ckpt\"\ncontrolnet_config: \"configs/sparsectrl/latent_condition.yaml\"\nadapter_lora_path: \"models/Motion_Module/v3_sd15_adapter.ckpt\"\n\ncfg_scale: 7.5 # in default realistic classifer-free guidance\nnegative_prompt: \"ugly, deformed, noisy, blurry, distorted, out of focus, bad anatomy, extra limbs, poorly drawn face, poorly drawn hands, missing fingers\"\n\ninference_steps: 100 # the total denosing step for inference\nguidance_scale: 0.3 # which scale of time step to end guidance\nguidance_steps: 40 # the step for guidance in inference, no more than 1000*guidance_scale, the remaining steps (inference_steps-guidance_steps) is performed without gudiance\nwarm_up_steps: 10\ncool_up_steps: 10\n\nmotion_guidance_weight: 2000\nmotion_guidance_blocks: ['up_blocks.1']\n\nadd_noise_step: 400" }, { "path": "configs/i2v_sketch.jsonl", "content": "{\"video_path\":\"reference_videos/sample_white_tiger.mp4\", \"condition_image_paths\":[\"condition_images/scribble/lion_forest.png\"], \"new_prompt\": \"Lion, walks in the forest\"}" }, { "path": "configs/i2v_sketch.yaml", "content": "motion_module: \"models/Motion_Module/v3_sd15_mm.ckpt\"\ndreambooth_path: \"models/DreamBooth_LoRA/realisticVisionV60B1_v51VAE.safetensors\"\nmodel_config: \"configs/model_config/model_config.yaml\"\ncontrolnet_config: \"configs/sparsectrl/image_condition.yaml\"\ncontrolnet_path: \"models/SparseCtrl/v3_sd15_sparsectrl_scribble.ckpt\"\nadapter_lora_path: \"models/Motion_Module/v3_sd15_adapter.ckpt\"\n\ncfg_scale: 7.5 # in default realistic classifer-free guidance\nnegative_prompt: \"ugly, deformed, noisy, blurry, distorted, out of focus, bad anatomy, extra limbs, poorly drawn face, poorly drawn hands, missing fingers\"\n\ninference_steps: 200 # the total denosing step for inference\nguidance_scale: 0.4 # which scale of time step to end guidance\nguidance_steps: 120 # the step for guidance in inference, no more than 1000*guidance_scale, the remaining steps (inference_steps-guidance_steps) is performed without gudiance\nwarm_up_steps: 10\ncool_up_steps: 10\n\nmotion_guidance_weight: 2000\nmotion_guidance_blocks: ['up_blocks.1']\n\nadd_noise_step: 400" }, { "path": "configs/model_config/inference-v1.yaml", "content": "unet_additional_kwargs:\n use_inflated_groupnorm: true # from config v3\n\n\n use_motion_module: true\n motion_module_resolutions: [1,2,4,8]\n motion_module_mid_block: false\n motion_module_decoder_only: false\n motion_module_type: \"Vanilla\"\n \n motion_module_kwargs:\n num_attention_heads: 8\n num_transformer_block: 1\n attention_block_types: [ \"Temporal_Self\", \"Temporal_Self\" ]\n temporal_position_encoding: true\n temporal_position_encoding_max_len: 32\n temporal_attention_dim_div: 1\n zero_initialize: true # from config v3\n\nnoise_scheduler_kwargs:\n beta_start: 0.00085\n beta_end: 0.012\n beta_schedule: \"linear\"\n steps_offset: 1\n clip_sample: False\n" }, { "path": "configs/model_config/inference-v2.yaml", "content": "unet_additional_kwargs:\n use_inflated_groupnorm: true\n unet_use_cross_frame_attention: false\n unet_use_temporal_attention: false\n use_motion_module: true\n motion_module_resolutions: [1,2,4,8]\n motion_module_mid_block: true\n motion_module_decoder_only: false\n motion_module_type: \"Vanilla\"\n\n motion_module_kwargs:\n num_attention_heads: 8\n num_transformer_block: 1\n attention_block_types: [ \"Temporal_Self\", \"Temporal_Self\" ]\n temporal_position_encoding: true\n temporal_position_encoding_max_len: 32\n temporal_attention_dim_div: 1\n\nnoise_scheduler_kwargs:\n beta_start: 0.00085\n beta_end: 0.012\n beta_schedule: \"linear\"\n steps_offset: 1\n clip_sample: False\n" }, { "path": "configs/model_config/inference-v3.yaml", "content": "unet_additional_kwargs:\n use_inflated_groupnorm: true\n use_motion_module: true\n motion_module_resolutions: [1,2,4,8]\n motion_module_mid_block: false\n motion_module_type: Vanilla\n\n motion_module_kwargs:\n num_attention_heads: 8\n num_transformer_block: 1\n attention_block_types: [ \"Temporal_Self\", \"Temporal_Self\" ]\n temporal_position_encoding: true\n temporal_position_encoding_max_len: 32\n temporal_attention_dim_div: 1\n zero_initialize: true\n\nnoise_scheduler_kwargs:\n beta_start: 0.00085\n beta_end: 0.012\n beta_schedule: \"linear\"\n steps_offset: 1\n clip_sample: False\n" }, { "path": "configs/model_config/model_config copy.yaml", "content": "unet_additional_kwargs:\n use_inflated_groupnorm: true # from config v3\n use_motion_module: true\n motion_module_resolutions: [1,2,4,8]\n motion_module_mid_block: false\n motion_module_type: \"Vanilla\"\n\n motion_module_kwargs:\n num_attention_heads: 8\n num_transformer_block: 1\n attention_block_types: [ \"Temporal_Self\", \"Temporal_Self\" ]\n temporal_position_encoding: true\n temporal_position_encoding_max_len: 32\n temporal_attention_dim_div: 1\n zero_initialize: true # from config v3\n\nnoise_scheduler_kwargs:\n beta_start: 0.00085\n beta_end: 0.012\n beta_schedule: \"linear\"\n steps_offset: 1\n clip_sample: False" }, { "path": "configs/model_config/model_config.yaml", "content": "unet_additional_kwargs:\n use_inflated_groupnorm: true\n use_motion_module: true\n motion_module_resolutions: [1,2,4,8]\n motion_module_mid_block: false\n motion_module_type: \"Vanilla\"\n\n motion_module_kwargs:\n num_attention_heads: 8\n num_transformer_block: 1\n attention_block_types: [ \"Temporal_Self\", \"Temporal_Self\" ]\n temporal_position_encoding: true\n temporal_attention_dim_div: 1\n zero_initialize: true\n\nnoise_scheduler_kwargs:\n beta_start: 0.00085\n beta_end: 0.012\n beta_schedule: \"linear\"\n steps_offset: 1\n clip_sample: false" }, { "path": "configs/model_config/model_config_public.yaml", "content": "unet_additional_kwargs:\n use_inflated_groupnorm: true # from config v3\n unet_use_cross_frame_attention: false\n unet_use_temporal_attention: false\n use_motion_module: true\n motion_module_resolutions: [1,2,4,8]\n motion_module_mid_block: false\n motion_module_decoder_only: false\n motion_module_type: \"Vanilla\"\n \n motion_module_kwargs:\n num_attention_heads: 8\n num_transformer_block: 1\n attention_block_types: [ \"Temporal_Self\", \"Temporal_Self\" ]\n temporal_position_encoding: true\n temporal_position_encoding_max_len: 32\n temporal_attention_dim_div: 1\n zero_initialize: true # from config v3\n\nnoise_scheduler_kwargs:\n beta_start: 0.00085\n beta_end: 0.012\n beta_schedule: \"linear\"\n steps_offset: 1\n clip_sample: False\n" }, { "path": "configs/sparsectrl/image_condition.yaml", "content": "controlnet_additional_kwargs:\n set_noisy_sample_input_to_zero: true\n use_simplified_condition_embedding: false\n conditioning_channels: 3\n\n use_motion_module: true\n motion_module_resolutions: [1,2,4,8]\n motion_module_mid_block: false\n motion_module_type: \"Vanilla\"\n\n motion_module_kwargs:\n num_attention_heads: 8\n num_transformer_block: 1\n attention_block_types: [ \"Temporal_Self\" ]\n temporal_position_encoding: true\n temporal_position_encoding_max_len: 32\n temporal_attention_dim_div: 1\n" }, { "path": "configs/sparsectrl/latent_condition.yaml", "content": "controlnet_additional_kwargs:\n set_noisy_sample_input_to_zero: true\n use_simplified_condition_embedding: true\n conditioning_channels: 4\n\n use_motion_module: true\n motion_module_resolutions: [1,2,4,8]\n motion_module_mid_block: false\n motion_module_type: \"Vanilla\"\n\n motion_module_kwargs:\n num_attention_heads: 8\n num_transformer_block: 1\n attention_block_types: [ \"Temporal_Self\" ]\n temporal_position_encoding: true\n temporal_position_encoding_max_len: 32\n temporal_attention_dim_div: 1\n" }, { "path": "configs/t2v_camera.jsonl", "content": "{\"video_path\":\"reference_videos/camera_zoom_in.mp4\", \"new_prompt\": \"Relics on the seabed\", \"seed\": 42}\n{\"video_path\":\"reference_videos/camera_zoom_in.mp4\", \"new_prompt\": \"A road in the mountain\", \"seed\": 42}\n{\"video_path\":\"reference_videos/camera_zoom_in.mp4\", \"new_prompt\": \"Caves, a path for exploration\", \"seed\": 2026}\n{\"video_path\":\"reference_videos/camera_zoom_in.mp4\", \"new_prompt\": \"Railway for train\"}\n{\"video_path\":\"reference_videos/camera_zoom_out.mp4\", \"new_prompt\": \"Tree, in the mountain\", \"seed\": 2026}\n{\"video_path\":\"reference_videos/camera_zoom_out.mp4\", \"new_prompt\": \"Red car on the track\", \"seed\": 2026}\n{\"video_path\":\"reference_videos/camera_zoom_out.mp4\", \"new_prompt\": \"Man, standing in his garden.\", \"seed\": 2026}\n{\"video_path\":\"reference_videos/camera_1.mp4\", \"new_prompt\": \"A island, on the ocean, sunny day\"}\n{\"video_path\":\"reference_videos/camera_1.mp4\", \"new_prompt\": \"A tower, with fireworks\"}\n{\"video_path\":\"reference_videos/camera_pan_up.mp4\", \"new_prompt\": \"Beautiful house, around with flowers\", \"seed\": 42}\n{\"video_path\":\"reference_videos/camera_translation_2.mp4\", \"new_prompt\": \"Forest, in winter\", \"seed\": 2028}\n{\"video_path\":\"reference_videos/camera_pan_down.mp4\", \"new_prompt\": \"Eagle, standing in the tree\", \"seed\": 2026}" }, { "path": "configs/t2v_camera.yaml", "content": "\nmotion_module: \"models/Motion_Module/v3_sd15_mm.ckpt\"\ndreambooth_path: \"models/DreamBooth_LoRA/realisticVisionV60B1_v51VAE.safetensors\"\nmodel_config: \"configs/model_config/model_config.yaml\"\n\ncfg_scale: 7.5 # in default realistic classifer-free guidance\nnegative_prompt: \"bad anatomy, extra limbs, ugly, deformed, noisy, blurry, distorted, out of focus, poorly drawn face, poorly drawn hands, missing fingers\"\npostive_prompt: \" 8k, high detailed, best quality, film grain, Fujifilm XT3\"\n\ninference_steps: 100 # the total denosing step for inference\nguidance_scale: 0.3 # which scale of time step to end guidance 0.2/40\nguidance_steps: 50 # the step for guidance in inference, no more than 1000*guidance_scale, the remaining steps (inference_steps-guidance_steps) is performed without gudiance\nwarm_up_steps: 10\ncool_up_steps: 10\n\nmotion_guidance_weight: 2000\nmotion_guidance_blocks: ['up_blocks.1']\n\nadd_noise_step: 400" }, { "path": "configs/t2v_object.jsonl", "content": "{\"video_path\":\"reference_videos/sample_astronaut.mp4\", \"new_prompt\": \"Robot, walks in the street.\",\"seed\":59}\n{\"video_path\":\"reference_videos/sample_cat.mp4\", \"new_prompt\": \"Tiger, raises its head.\", \"seed\": 2025}\n{\"video_path\":\"reference_videos/sample_leaves.mp4\", \"new_prompt\": \"Petals falling in the wind.\",\"seed\":3407}\n{\"video_path\":\"reference_videos/sample_fox.mp4\", \"new_prompt\": \"Cat, turns its head in the living room.\"}\n{\"video_path\":\"reference_videos/sample_blackswan.mp4\", \"new_prompt\": \"Duck, swims in the river.\",\"seed\":3407}\n{\"video_path\":\"reference_videos/sample_cow.mp4\", \"new_prompt\": \"Pig, drinks water on beach.\",\"seed\":3407}" }, { "path": "configs/t2v_object.yaml", "content": "\nmotion_module: \"models/Motion_Module/v3_sd15_mm.ckpt\"\ndreambooth_path: \"models/DreamBooth_LoRA/realisticVisionV60B1_v51VAE.safetensors\"\nmodel_config: \"configs/model_config/model_config.yaml\"\n\ncfg_scale: 7.5 # in default realistic classifer-free guidance\nnegative_prompt: \"bad anatomy, extra limbs, ugly, deformed, noisy, blurry, distorted, out of focus, poorly drawn face, poorly drawn hands, missing fingers\"\npostive_prompt: \"8k, high detailed, best quality, film grain, Fujifilm XT3\"\n\ninference_steps: 300 # the total denosing step for inference\nguidance_scale: 0.4 # which scale of time step to end guidance\nguidance_steps: 180 # the step for guidance in inference, no more than 1000*guidance_scale, the remaining steps (inference_steps-guidance_steps) is performed without gudiance\nwarm_up_steps: 10\ncool_up_steps: 10\n\nmotion_guidance_weight: 2000\nmotion_guidance_blocks: ['up_blocks.1',]\n\nadd_noise_step: 400" }, { "path": "environment.yaml", "content": "name: motionclone\nchannels:\n - pytorch\n - nvidia\ndependencies:\n - python=3.11.3\n - pytorch=2.0.1\n - torchvision=0.15.2\n - pytorch-cuda=11.8\n - pip\n - pip:\n - accelerate\n - diffusers==0.16.0\n - transformers==4.28.1\n - xformers==0.0.20\n - imageio[ffmpeg]\n - decord==0.6.0\n - gdown\n - einops\n - omegaconf\n - safetensors\n - gradio\n - wandb\n - triton\n - opencv-python" }, { "path": "generated_videos/inference_config.json", "content": "motion_module: models/Motion_Module/v3_sd15_mm.ckpt\ndreambooth_path: models/DreamBooth_LoRA/realisticVisionV60B1_v51VAE.safetensors\nmodel_config: configs/model_config/model_config.yaml\ncontrolnet_config: configs/sparsectrl/image_condition.yaml\ncontrolnet_path: models/SparseCtrl/v3_sd15_sparsectrl_scribble.ckpt\nadapter_lora_path: models/Motion_Module/v3_sd15_adapter.ckpt\ncfg_scale: 7.5\nnegative_prompt: ugly, deformed, noisy, blurry, distorted, out of focus, bad anatomy,\n extra limbs, poorly drawn face, poorly drawn hands, missing fingers\ninference_steps: 200\nguidance_scale: 0.4\nguidance_steps: 120\nwarm_up_steps: 10\ncool_up_steps: 10\nmotion_guidance_weight: 2000\nmotion_guidance_blocks:\n- up_blocks.1\nadd_noise_step: 400\nwidth: 512\nheight: 512\nvideo_length: 16\n" }, { "path": "i2v_video_sample.py", "content": "import argparse\nfrom omegaconf import OmegaConf\nimport torch\nfrom diffusers import AutoencoderKL, DDIMScheduler\nfrom transformers import CLIPTextModel, CLIPTokenizer\nfrom motionclone.models.unet import UNet3DConditionModel\nfrom motionclone.models.sparse_controlnet import SparseControlNetModel\nfrom motionclone.pipelines.pipeline_animation import AnimationPipeline\nfrom motionclone.utils.util import load_weights, auto_download\nfrom diffusers.utils.import_utils import is_xformers_available\nfrom motionclone.utils.motionclone_functions import *\nimport json\nfrom motionclone.utils.xformer_attention import *\n\n\ndef main(args):\n\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = args.visible_gpu or str(os.getenv('CUDA_VISIBLE_DEVICES', 0))\n \n config = OmegaConf.load(args.inference_config)\n adopted_dtype = torch.float16\n device = \"cuda\"\n set_all_seed(42)\n \n tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_path, subfolder=\"tokenizer\")\n text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_path, subfolder=\"text_encoder\").to(device).to(dtype=adopted_dtype)\n vae = AutoencoderKL.from_pretrained(args.pretrained_model_path, subfolder=\"vae\").to(device).to(dtype=adopted_dtype)\n \n config.width = config.get(\"W\", args.W)\n config.height = config.get(\"H\", args.H)\n config.video_length = config.get(\"L\", args.L)\n \n if not os.path.exists(args.generated_videos_save_dir):\n os.makedirs(args.generated_videos_save_dir)\n OmegaConf.save(config, os.path.join(args.generated_videos_save_dir,\"inference_config.json\"))\n \n model_config = OmegaConf.load(config.get(\"model_config\", \"\"))\n unet = UNet3DConditionModel.from_pretrained_2d(args.pretrained_model_path, subfolder=\"unet\", unet_additional_kwargs=OmegaConf.to_container(model_config.unet_additional_kwargs),).to(device).to(dtype=adopted_dtype)\n \n # load controlnet model\n controlnet = None\n if config.get(\"controlnet_path\", \"\") != \"\":\n # assert model_config.get(\"controlnet_images\", \"\") != \"\"\n assert config.get(\"controlnet_config\", \"\") != \"\"\n \n unet.config.num_attention_heads = 8\n unet.config.projection_class_embeddings_input_dim = None\n\n controlnet_config = OmegaConf.load(config.controlnet_config)\n controlnet = SparseControlNetModel.from_unet(unet, controlnet_additional_kwargs=controlnet_config.get(\"controlnet_additional_kwargs\", {})).to(device).to(dtype=adopted_dtype)\n\n auto_download(config.controlnet_path, is_dreambooth_lora=False)\n print(f\"loading controlnet checkpoint from {config.controlnet_path} ...\")\n controlnet_state_dict = torch.load(config.controlnet_path, map_location=\"cpu\")\n controlnet_state_dict = controlnet_state_dict[\"controlnet\"] if \"controlnet\" in controlnet_state_dict else controlnet_state_dict\n controlnet_state_dict = {name: param for name, param in controlnet_state_dict.items() if \"pos_encoder.pe\" not in name}\n controlnet_state_dict.pop(\"animatediff_config\", \"\")\n controlnet.load_state_dict(controlnet_state_dict)\n del controlnet_state_dict\n\n # set xformers\n if is_xformers_available() and (not args.without_xformers):\n unet.enable_xformers_memory_efficient_attention()\n\n pipeline = AnimationPipeline(\n vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet,\n controlnet=controlnet,\n scheduler=DDIMScheduler(**OmegaConf.to_container(model_config.noise_scheduler_kwargs)),\n ).to(device)\n \n pipeline = load_weights(\n pipeline,\n # motion module\n motion_module_path = config.get(\"motion_module\", \"\"),\n # domain adapter\n adapter_lora_path = config.get(\"adapter_lora_path\", \"\"),\n adapter_lora_scale = config.get(\"adapter_lora_scale\", 1.0),\n # image layer\n dreambooth_model_path = config.get(\"dreambooth_path\", \"\"),\n ).to(device)\n pipeline.text_encoder.to(dtype=adopted_dtype)\n \n # customized functions in motionclone_functions\n pipeline.scheduler.customized_step = schedule_customized_step.__get__(pipeline.scheduler)\n pipeline.scheduler.customized_set_timesteps = schedule_set_timesteps.__get__(pipeline.scheduler)\n pipeline.unet.forward = unet_customized_forward.__get__(pipeline.unet)\n pipeline.sample_video = sample_video.__get__(pipeline)\n pipeline.single_step_video = single_step_video.__get__(pipeline)\n pipeline.get_temp_attn_prob = get_temp_attn_prob.__get__(pipeline)\n pipeline.add_noise = add_noise.__get__(pipeline)\n pipeline.compute_temp_loss = compute_temp_loss.__get__(pipeline)\n pipeline.obtain_motion_representation = obtain_motion_representation.__get__(pipeline)\n \n for param in pipeline.unet.parameters():\n param.requires_grad = False\n for param in pipeline.controlnet.parameters():\n param.requires_grad = False\n \n pipeline.input_config, pipeline.unet.input_config = config, config\n pipeline.unet = prep_unet_attention(pipeline.unet,pipeline.input_config.motion_guidance_blocks)\n pipeline.unet = prep_unet_conv(pipeline.unet)\n pipeline.scheduler.customized_set_timesteps(config.inference_steps, config.guidance_steps,config.guidance_scale,device=device,timestep_spacing_type = \"uneven\")\n \n with open(args.examples, 'r') as files:\n for line in files:\n # prepare infor of each case\n example_infor = json.loads(line)\n config.video_path = example_infor[\"video_path\"]\n config.condition_image_path_list = example_infor[\"condition_image_paths\"]\n config.image_index = example_infor.get(\"image_index\",[0])\n assert len(config.image_index) == len(config.condition_image_path_list)\n config.new_prompt = example_infor[\"new_prompt\"] + config.get(\"positive_prompt\", \"\")\n config.controlnet_scale = example_infor.get(\"controlnet_scale\", 1.0)\n pipeline.input_config, pipeline.unet.input_config = config, config # update config\n \n # perform motion representation extraction\n seed_motion = seed_motion = example_infor.get(\"seed\", args.default_seed) \n generator = torch.Generator(device=pipeline.device)\n generator.manual_seed(seed_motion)\n if not os.path.exists(args.motion_representation_save_dir):\n os.makedirs(args.motion_representation_save_dir)\n motion_representation_path = os.path.join(args.motion_representation_save_dir, os.path.splitext(os.path.basename(config.video_path))[0] + '.pt') \n pipeline.obtain_motion_representation(generator= generator, motion_representation_path = motion_representation_path, use_controlnet=True,) \n \n # perform video generation\n seed = seed_motion # can assign other seed here\n generator = torch.Generator(device=pipeline.device)\n generator.manual_seed(seed)\n pipeline.input_config.seed = seed\n videos = pipeline.sample_video(generator = generator, add_controlnet=True,)\n\n videos = rearrange(videos, \"b c f h w -> b f h w c\")\n save_path = os.path.join(args.generated_videos_save_dir, os.path.splitext(os.path.basename(config.video_path))[0]\n + \"_\" + config.new_prompt.strip().replace(' ', '_') + str(seed_motion) + \"_\" +str(seed)+'.mp4') \n videos_uint8 = (videos[0] * 255).astype(np.uint8)\n imageio.mimwrite(save_path, videos_uint8, fps=8)\n print(save_path,\"is done\")\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--pretrained-model-path\", type=str, default=\"models/StableDiffusion\",)\n \n parser.add_argument(\"--inference_config\", type=str, default=\"configs/i2v_sketch.yaml\")\n parser.add_argument(\"--examples\", type=str, default=\"configs/i2v_sketch.jsonl\")\n parser.add_argument(\"--motion-representation-save-dir\", type=str, default=\"motion_representation/\")\n parser.add_argument(\"--generated-videos-save-dir\", type=str, default=\"generated_videos/\")\n \n parser.add_argument(\"--visible_gpu\", type=str, default=None)\n parser.add_argument(\"--default-seed\", type=int, default=76739)\n parser.add_argument(\"--L\", type=int, default=16)\n parser.add_argument(\"--W\", type=int, default=512)\n parser.add_argument(\"--H\", type=int, default=512)\n\n parser.add_argument(\"--without-xformers\", action=\"store_true\")\n\n args = parser.parse_args()\n main(args)\n" }, { "path": "models/Motion_Module/Put motion module checkpoints here.txt", "content": "" }, { "path": "motionclone/models/attention.py", "content": "# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py\n\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\n\nfrom diffusers.configuration_utils import ConfigMixin, register_to_config\nfrom diffusers.models.modeling_utils import ModelMixin\nfrom diffusers.utils import BaseOutput\nfrom diffusers.utils.import_utils import is_xformers_available\nfrom diffusers.models.attention import FeedForward, AdaLayerNorm\n\nfrom einops import rearrange, repeat\nimport pdb\n\n@dataclass\nclass Transformer3DModelOutput(BaseOutput):\n sample: torch.FloatTensor\n\n\nif is_xformers_available():\n import xformers\n import xformers.ops\nelse:\n xformers = None\n\n\nclass Transformer3DModel(ModelMixin, ConfigMixin):\n @register_to_config\n def __init__(\n self,\n num_attention_heads: int = 16,\n attention_head_dim: int = 88,\n in_channels: Optional[int] = None,\n num_layers: int = 1,\n dropout: float = 0.0,\n norm_num_groups: int = 32,\n cross_attention_dim: Optional[int] = None,\n attention_bias: bool = False,\n activation_fn: str = \"geglu\",\n num_embeds_ada_norm: Optional[int] = None,\n use_linear_projection: bool = False,\n only_cross_attention: bool = False,\n upcast_attention: bool = False,\n\n unet_use_cross_frame_attention=None,\n unet_use_temporal_attention=None,\n ):\n super().__init__()\n self.use_linear_projection = use_linear_projection\n self.num_attention_heads = num_attention_heads\n self.attention_head_dim = attention_head_dim\n inner_dim = num_attention_heads * attention_head_dim\n\n # Define input layers\n self.in_channels = in_channels\n\n self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)\n if use_linear_projection:\n self.proj_in = nn.Linear(in_channels, inner_dim)\n else:\n self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)\n\n # Define transformers blocks\n self.transformer_blocks = nn.ModuleList(\n [\n BasicTransformerBlock(\n inner_dim,\n num_attention_heads,\n attention_head_dim,\n dropout=dropout,\n cross_attention_dim=cross_attention_dim,\n activation_fn=activation_fn,\n num_embeds_ada_norm=num_embeds_ada_norm,\n attention_bias=attention_bias,\n only_cross_attention=only_cross_attention,\n upcast_attention=upcast_attention,\n\n unet_use_cross_frame_attention=unet_use_cross_frame_attention,\n unet_use_temporal_attention=unet_use_temporal_attention,\n )\n for d in range(num_layers)\n ]\n )\n\n # 4. Define output layers\n if use_linear_projection:\n self.proj_out = nn.Linear(in_channels, inner_dim)\n else:\n self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)\n\n def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, return_dict: bool = True):\n # Input\n assert hidden_states.dim() == 5, f\"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}.\"\n video_length = hidden_states.shape[2]\n hidden_states = rearrange(hidden_states, \"b c f h w -> (b f) c h w\")\n encoder_hidden_states = repeat(encoder_hidden_states, 'b n c -> (b f) n c', f=video_length)\n\n batch, channel, height, weight = hidden_states.shape\n residual = hidden_states\n\n hidden_states = self.norm(hidden_states)\n if not self.use_linear_projection:\n hidden_states = self.proj_in(hidden_states)\n inner_dim = hidden_states.shape[1]\n hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)\n else:\n inner_dim = hidden_states.shape[1]\n hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)\n hidden_states = self.proj_in(hidden_states)\n\n # Blocks\n for block in self.transformer_blocks:\n hidden_states = block(\n hidden_states,\n encoder_hidden_states=encoder_hidden_states,\n timestep=timestep,\n video_length=video_length\n )\n\n # Output\n if not self.use_linear_projection:\n hidden_states = (\n hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()\n )\n hidden_states = self.proj_out(hidden_states)\n else:\n hidden_states = self.proj_out(hidden_states)\n hidden_states = (\n hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()\n )\n\n output = hidden_states + residual\n\n output = rearrange(output, \"(b f) c h w -> b c f h w\", f=video_length)\n if not return_dict:\n return (output,)\n\n return Transformer3DModelOutput(sample=output)\n\n\nclass BasicTransformerBlock(nn.Module):\n def __init__(\n self,\n dim: int,\n num_attention_heads: int,\n attention_head_dim: int,\n dropout=0.0,\n cross_attention_dim: Optional[int] = None,\n activation_fn: str = \"geglu\",\n num_embeds_ada_norm: Optional[int] = None,\n attention_bias: bool = False,\n only_cross_attention: bool = False,\n upcast_attention: bool = False,\n\n unet_use_cross_frame_attention = None,\n unet_use_temporal_attention = None,\n ):\n super().__init__()\n self.only_cross_attention = only_cross_attention\n self.use_ada_layer_norm = num_embeds_ada_norm is not None\n self.unet_use_cross_frame_attention = unet_use_cross_frame_attention\n self.unet_use_temporal_attention = unet_use_temporal_attention\n\n # SC-Attn\n assert unet_use_cross_frame_attention is not None\n if unet_use_cross_frame_attention:\n self.attn1 = SparseCausalAttention2D(\n query_dim=dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n cross_attention_dim=cross_attention_dim if only_cross_attention else None,\n upcast_attention=upcast_attention,\n )\n else:\n self.attn1 = CrossAttention(\n query_dim=dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n upcast_attention=upcast_attention,\n )\n self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)\n\n # Cross-Attn\n if cross_attention_dim is not None:\n self.attn2 = CrossAttention(\n query_dim=dim,\n cross_attention_dim=cross_attention_dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n upcast_attention=upcast_attention,\n )\n else:\n self.attn2 = None\n\n if cross_attention_dim is not None:\n self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)\n else:\n self.norm2 = None\n\n # Feed-forward\n self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)\n self.norm3 = nn.LayerNorm(dim)\n\n # Temp-Attn\n assert unet_use_temporal_attention is not None\n if unet_use_temporal_attention:\n self.attn_temp = CrossAttention(\n query_dim=dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n upcast_attention=upcast_attention,\n )\n nn.init.zeros_(self.attn_temp.to_out[0].weight.data)\n self.norm_temp = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)\n\n def set_use_memory_efficient_attention_xformers(self, use_memory_efficient_attention_xformers: bool, op=None):\n if not is_xformers_available():\n print(\"Here is how to install it\")\n raise ModuleNotFoundError(\n \"Refer to https://github.com/facebookresearch/xformers for more information on how to install\"\n \" xformers\",\n name=\"xformers\",\n )\n elif not torch.cuda.is_available():\n raise ValueError(\n \"torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is only\"\n \" available for GPU \"\n )\n else:\n try:\n # Make sure we can run the memory efficient attention\n _ = xformers.ops.memory_efficient_attention(\n torch.randn((1, 2, 40), device=\"cuda\"),\n torch.randn((1, 2, 40), device=\"cuda\"),\n torch.randn((1, 2, 40), device=\"cuda\"),\n )\n except Exception as e:\n raise e\n self.attn1._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers\n if self.attn2 is not None:\n self.attn2._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers\n # self.attn_temp._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers\n\n def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, attention_mask=None, video_length=None):\n # SparseCausal-Attention\n norm_hidden_states = (\n self.norm1(hidden_states, timestep) if self.use_ada_layer_norm else self.norm1(hidden_states)\n )\n\n # if self.only_cross_attention:\n # hidden_states = (\n # self.attn1(norm_hidden_states, encoder_hidden_states, attention_mask=attention_mask) + hidden_states\n # )\n # else:\n # hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, video_length=video_length) + hidden_states\n\n # pdb.set_trace()\n if self.unet_use_cross_frame_attention:\n hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, video_length=video_length) + hidden_states\n else:\n hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask) + hidden_states\n\n if self.attn2 is not None:\n # Cross-Attention\n norm_hidden_states = (\n self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)\n )\n hidden_states = (\n self.attn2(\n norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask\n )\n + hidden_states\n )\n\n # Feed-forward\n hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states\n\n # Temporal-Attention\n if self.unet_use_temporal_attention:\n d = hidden_states.shape[1]\n hidden_states = rearrange(hidden_states, \"(b f) d c -> (b d) f c\", f=video_length)\n norm_hidden_states = (\n self.norm_temp(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_temp(hidden_states)\n )\n hidden_states = self.attn_temp(norm_hidden_states) + hidden_states\n hidden_states = rearrange(hidden_states, \"(b d) f c -> (b f) d c\", d=d)\n\n return hidden_states\n\nclass CrossAttention(nn.Module):\n r\"\"\"\n A cross attention layer.\n\n Parameters:\n query_dim (`int`): The number of channels in the query.\n cross_attention_dim (`int`, *optional*):\n The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.\n heads (`int`, *optional*, defaults to 8): The number of heads to use for multi-head attention.\n dim_head (`int`, *optional*, defaults to 64): The number of channels in each head.\n dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.\n bias (`bool`, *optional*, defaults to False):\n Set to `True` for the query, key, and value linear layers to contain a bias parameter.\n \"\"\"\n\n def __init__(\n self,\n query_dim: int,\n cross_attention_dim: Optional[int] = None,\n heads: int = 8,\n dim_head: int = 64,\n dropout: float = 0.0,\n bias=False,\n upcast_attention: bool = False,\n upcast_softmax: bool = False,\n added_kv_proj_dim: Optional[int] = None,\n norm_num_groups: Optional[int] = None,\n ):\n super().__init__()\n inner_dim = dim_head * heads\n cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim\n self.upcast_attention = upcast_attention\n self.upcast_softmax = upcast_softmax\n\n self.scale = dim_head**-0.5\n\n self.heads = heads\n # for slice_size > 0 the attention score computation\n # is split across the batch axis to save memory\n # You can set slice_size with `set_attention_slice`\n self.sliceable_head_dim = heads\n self._slice_size = None\n self._use_memory_efficient_attention_xformers = False\n self.added_kv_proj_dim = added_kv_proj_dim\n\n #### add processer\n self.processor = None\n\n if norm_num_groups is not None:\n self.group_norm = nn.GroupNorm(num_channels=inner_dim, num_groups=norm_num_groups, eps=1e-5, affine=True)\n else:\n self.group_norm = None\n\n self.to_q = nn.Linear(query_dim, inner_dim, bias=bias)\n self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias)\n self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias)\n\n if self.added_kv_proj_dim is not None:\n self.add_k_proj = nn.Linear(added_kv_proj_dim, cross_attention_dim)\n self.add_v_proj = nn.Linear(added_kv_proj_dim, cross_attention_dim)\n\n self.to_out = nn.ModuleList([])\n self.to_out.append(nn.Linear(inner_dim, query_dim))\n self.to_out.append(nn.Dropout(dropout))\n\n def reshape_heads_to_batch_dim(self, tensor):\n batch_size, seq_len, dim = tensor.shape\n head_size = self.heads\n tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)\n tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size * head_size, seq_len, dim // head_size)\n return tensor\n\n def reshape_batch_dim_to_heads(self, tensor):\n batch_size, seq_len, dim = tensor.shape\n head_size = self.heads\n tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)\n tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)\n return tensor\n\n def set_attention_slice(self, slice_size):\n if slice_size is not None and slice_size > self.sliceable_head_dim:\n raise ValueError(f\"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.\")\n\n self._slice_size = slice_size\n\n def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):\n batch_size, sequence_length, _ = hidden_states.shape\n\n encoder_hidden_states = encoder_hidden_states\n\n if self.group_norm is not None:\n hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = self.to_q(hidden_states)\n dim = query.shape[-1]\n # query = self.reshape_heads_to_batch_dim(query) # move backwards\n\n if self.added_kv_proj_dim is not None:\n key = self.to_k(hidden_states)\n value = self.to_v(hidden_states)\n encoder_hidden_states_key_proj = self.add_k_proj(encoder_hidden_states)\n encoder_hidden_states_value_proj = self.add_v_proj(encoder_hidden_states)\n\n ######record###### record before reshape heads to batch dim\n if self.processor is not None:\n self.processor.record_qkv(self, hidden_states, query, key, value, attention_mask)\n ##################\n\n key = self.reshape_heads_to_batch_dim(key)\n value = self.reshape_heads_to_batch_dim(value)\n encoder_hidden_states_key_proj = self.reshape_heads_to_batch_dim(encoder_hidden_states_key_proj)\n encoder_hidden_states_value_proj = self.reshape_heads_to_batch_dim(encoder_hidden_states_value_proj)\n\n key = torch.concat([encoder_hidden_states_key_proj, key], dim=1)\n value = torch.concat([encoder_hidden_states_value_proj, value], dim=1)\n else:\n encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states\n key = self.to_k(encoder_hidden_states)\n value = self.to_v(encoder_hidden_states)\n\n ######record######\n if self.processor is not None:\n self.processor.record_qkv(self, hidden_states, query, key, value, attention_mask)\n ##################\n\n key = self.reshape_heads_to_batch_dim(key)\n value = self.reshape_heads_to_batch_dim(value)\n\n query = self.reshape_heads_to_batch_dim(query) # reshape query\n\n if attention_mask is not None:\n if attention_mask.shape[-1] != query.shape[1]:\n target_length = query.shape[1]\n attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)\n attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)\n\n ######record######\n if self.processor is not None:\n self.processor.record_attn_mask(self, hidden_states, query, key, value, attention_mask)\n ##################\n\n # attention, what we cannot get enough of\n if self._use_memory_efficient_attention_xformers:\n hidden_states = self._memory_efficient_attention_xformers(query, key, value, attention_mask)\n # Some versions of xformers return output in fp32, cast it back to the dtype of the input\n hidden_states = hidden_states.to(query.dtype)\n else:\n if self._slice_size is None or query.shape[0] // self._slice_size == 1:\n hidden_states = self._attention(query, key, value, attention_mask)\n else:\n hidden_states = self._sliced_attention(query, key, value, sequence_length, dim, attention_mask)\n\n # linear proj\n hidden_states = self.to_out[0](hidden_states)\n\n # dropout\n hidden_states = self.to_out[1](hidden_states)\n return hidden_states\n\n def _attention(self, query, key, value, attention_mask=None):\n if self.upcast_attention:\n query = query.float()\n key = key.float()\n\n attention_scores = torch.baddbmm(\n torch.empty(query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device),\n query,\n key.transpose(-1, -2),\n beta=0,\n alpha=self.scale,\n )\n\n if attention_mask is not None:\n attention_scores = attention_scores + attention_mask\n\n if self.upcast_softmax:\n attention_scores = attention_scores.float()\n\n attention_probs = attention_scores.softmax(dim=-1)\n\n # cast back to the original dtype\n attention_probs = attention_probs.to(value.dtype)\n\n # compute attention output\n hidden_states = torch.bmm(attention_probs, value)\n\n # reshape hidden_states\n hidden_states = self.reshape_batch_dim_to_heads(hidden_states)\n return hidden_states\n\n def _sliced_attention(self, query, key, value, sequence_length, dim, attention_mask):\n batch_size_attention = query.shape[0]\n hidden_states = torch.zeros(\n (batch_size_attention, sequence_length, dim // self.heads), device=query.device, dtype=query.dtype\n )\n slice_size = self._slice_size if self._slice_size is not None else hidden_states.shape[0]\n for i in range(hidden_states.shape[0] // slice_size):\n start_idx = i * slice_size\n end_idx = (i + 1) * slice_size\n\n query_slice = query[start_idx:end_idx]\n key_slice = key[start_idx:end_idx]\n\n if self.upcast_attention:\n query_slice = query_slice.float()\n key_slice = key_slice.float()\n\n attn_slice = torch.baddbmm(\n torch.empty(slice_size, query.shape[1], key.shape[1], dtype=query_slice.dtype, device=query.device),\n query_slice,\n key_slice.transpose(-1, -2),\n beta=0,\n alpha=self.scale,\n )\n\n if attention_mask is not None:\n attn_slice = attn_slice + attention_mask[start_idx:end_idx]\n\n if self.upcast_softmax:\n attn_slice = attn_slice.float()\n\n attn_slice = attn_slice.softmax(dim=-1)\n\n # cast back to the original dtype\n attn_slice = attn_slice.to(value.dtype)\n attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])\n\n hidden_states[start_idx:end_idx] = attn_slice\n\n # reshape hidden_states\n hidden_states = self.reshape_batch_dim_to_heads(hidden_states)\n return hidden_states\n\n def _memory_efficient_attention_xformers(self, query, key, value, attention_mask):\n # TODO attention_mask\n query = query.contiguous()\n key = key.contiguous()\n value = value.contiguous()\n hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)\n hidden_states = self.reshape_batch_dim_to_heads(hidden_states)\n return hidden_states\n \n def set_processor(self, processor: \"AttnProcessor\") -> None:\n r\"\"\"\n Set the attention processor to use.\n\n Args:\n processor (`AttnProcessor`):\n The attention processor to use.\n \"\"\"\n # if current processor is in `self._modules` and if passed `processor` is not, we need to\n # pop `processor` from `self._modules`\n if (\n hasattr(self, \"processor\")\n and isinstance(self.processor, torch.nn.Module)\n and not isinstance(processor, torch.nn.Module)\n ):\n logger.info(f\"You are removing possibly trained weights of {self.processor} with {processor}\")\n self._modules.pop(\"processor\")\n\n self.processor = processor\n \n def get_attention_scores(\n self, query: torch.Tensor, key: torch.Tensor, attention_mask: torch.Tensor = None\n ) -> torch.Tensor:\n r\"\"\"\n Compute the attention scores.\n\n Args:\n query (`torch.Tensor`): The query tensor.\n key (`torch.Tensor`): The key tensor.\n attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.\n\n Returns:\n `torch.Tensor`: The attention probabilities/scores.\n \"\"\"\n dtype = query.dtype\n if self.upcast_attention:\n query = query.float()\n key = key.float()\n\n if attention_mask is None:\n baddbmm_input = torch.empty(\n query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device\n )\n beta = 0\n else:\n baddbmm_input = attention_mask\n beta = 1\n\n \n\n attention_scores = torch.baddbmm(\n baddbmm_input,\n query,\n key.transpose(-1, -2),\n beta=beta,\n alpha=self.scale,\n )\n del baddbmm_input\n\n if self.upcast_softmax:\n attention_scores = attention_scores.float()\n\n attention_probs = attention_scores.softmax(dim=-1)\n del attention_scores\n\n attention_probs = attention_probs.to(dtype)\n\n return attention_probs\n" }, { "path": "motionclone/models/motion_module.py", "content": "from dataclasses import dataclass\nfrom typing import List, Optional, Tuple, Union\n\nimport torch\nimport numpy as np\nimport torch.nn.functional as F\nfrom torch import nn\nimport torchvision\n\nfrom diffusers.configuration_utils import ConfigMixin, register_to_config\nfrom diffusers.models.modeling_utils import ModelMixin\nfrom diffusers.utils import BaseOutput\nfrom diffusers.utils.import_utils import is_xformers_available\nfrom diffusers.models.attention import FeedForward\nfrom .attention import CrossAttention\n\nfrom einops import rearrange, repeat\nimport math\n\n\ndef zero_module(module):\n # Zero out the parameters of a module and return it.\n for p in module.parameters():\n p.detach().zero_()\n return module\n\n\n@dataclass\nclass TemporalTransformer3DModelOutput(BaseOutput):\n sample: torch.FloatTensor\n\n\nif is_xformers_available():\n import xformers\n import xformers.ops\nelse:\n xformers = None\n\n\ndef get_motion_module( # 只能返回VanillaTemporalModule类\n in_channels,\n motion_module_type: str, \n motion_module_kwargs: dict\n):\n if motion_module_type == \"Vanilla\":\n return VanillaTemporalModule(in_channels=in_channels, **motion_module_kwargs,) \n else:\n raise ValueError\n\n\nclass VanillaTemporalModule(nn.Module):\n def __init__(\n self,\n in_channels,\n num_attention_heads = 8,\n num_transformer_block = 2,\n attention_block_types =( \"Temporal_Self\", \"Temporal_Self\" ),\n cross_frame_attention_mode = None,\n temporal_position_encoding = False,\n temporal_position_encoding_max_len = 32,\n temporal_attention_dim_div = 1,\n zero_initialize = True,\n ):\n super().__init__()\n \n self.temporal_transformer = TemporalTransformer3DModel(\n in_channels=in_channels,\n num_attention_heads=num_attention_heads,\n attention_head_dim=in_channels // num_attention_heads // temporal_attention_dim_div,\n num_layers=num_transformer_block,\n attention_block_types=attention_block_types,\n cross_frame_attention_mode=cross_frame_attention_mode,\n temporal_position_encoding=temporal_position_encoding,\n temporal_position_encoding_max_len=temporal_position_encoding_max_len,\n )\n \n if zero_initialize:\n self.temporal_transformer.proj_out = zero_module(self.temporal_transformer.proj_out)\n\n def forward(self, input_tensor, temb, encoder_hidden_states, attention_mask=None, anchor_frame_idx=None):\n hidden_states = input_tensor\n hidden_states = self.temporal_transformer(hidden_states, encoder_hidden_states, attention_mask)\n\n output = hidden_states\n return output\n\n\nclass TemporalTransformer3DModel(nn.Module):\n def __init__(\n self,\n in_channels,\n num_attention_heads,\n attention_head_dim,\n\n num_layers,\n attention_block_types = ( \"Temporal_Self\", \"Temporal_Self\", ), # 两个TempAttn \n dropout = 0.0,\n norm_num_groups = 32,\n cross_attention_dim = 768,\n activation_fn = \"geglu\",\n attention_bias = False,\n upcast_attention = False,\n \n cross_frame_attention_mode = None,\n temporal_position_encoding = False,\n temporal_position_encoding_max_len = 24,\n ):\n super().__init__()\n\n inner_dim = num_attention_heads * attention_head_dim\n\n self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)\n self.proj_in = nn.Linear(in_channels, inner_dim)\n\n self.transformer_blocks = nn.ModuleList(\n [\n TemporalTransformerBlock(\n dim=inner_dim,\n num_attention_heads=num_attention_heads,\n attention_head_dim=attention_head_dim,\n attention_block_types=attention_block_types,\n dropout=dropout,\n norm_num_groups=norm_num_groups,\n cross_attention_dim=cross_attention_dim,\n activation_fn=activation_fn,\n attention_bias=attention_bias,\n upcast_attention=upcast_attention,\n cross_frame_attention_mode=cross_frame_attention_mode,\n temporal_position_encoding=temporal_position_encoding,\n temporal_position_encoding_max_len=temporal_position_encoding_max_len,\n )\n for d in range(num_layers)\n ]\n )\n self.proj_out = nn.Linear(inner_dim, in_channels) \n \n def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):\n assert hidden_states.dim() == 5, f\"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}.\"\n video_length = hidden_states.shape[2]\n hidden_states = rearrange(hidden_states, \"b c f h w -> (b f) c h w\")\n\n batch, channel, height, weight = hidden_states.shape\n residual = hidden_states\n\n hidden_states = self.norm(hidden_states)\n inner_dim = hidden_states.shape[1]\n hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)\n hidden_states = self.proj_in(hidden_states)\n\n # Transformer Blocks\n for block in self.transformer_blocks:\n hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states, video_length=video_length)\n \n # output\n hidden_states = self.proj_out(hidden_states)\n hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()\n\n output = hidden_states + residual\n output = rearrange(output, \"(b f) c h w -> b c f h w\", f=video_length)\n \n return output\n\n\nclass TemporalTransformerBlock(nn.Module):\n def __init__(\n self,\n dim,\n num_attention_heads,\n attention_head_dim,\n attention_block_types = ( \"Temporal_Self\", \"Temporal_Self\", ),\n dropout = 0.0,\n norm_num_groups = 32,\n cross_attention_dim = 768,\n activation_fn = \"geglu\",\n attention_bias = False,\n upcast_attention = False,\n cross_frame_attention_mode = None,\n temporal_position_encoding = False,\n temporal_position_encoding_max_len = 24,\n ):\n super().__init__()\n\n attention_blocks = []\n norms = []\n \n for block_name in attention_block_types:\n attention_blocks.append(\n VersatileAttention(\n attention_mode=block_name.split(\"_\")[0],\n cross_attention_dim=cross_attention_dim if block_name.endswith(\"_Cross\") else None,\n \n query_dim=dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n upcast_attention=upcast_attention,\n \n cross_frame_attention_mode=cross_frame_attention_mode,\n temporal_position_encoding=temporal_position_encoding,\n temporal_position_encoding_max_len=temporal_position_encoding_max_len,\n )\n )\n norms.append(nn.LayerNorm(dim))\n \n self.attention_blocks = nn.ModuleList(attention_blocks)\n self.norms = nn.ModuleList(norms)\n\n self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)\n self.ff_norm = nn.LayerNorm(dim)\n\n\n def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None):\n for attention_block, norm in zip(self.attention_blocks, self.norms):\n norm_hidden_states = norm(hidden_states)\n hidden_states = attention_block(\n norm_hidden_states,\n encoder_hidden_states=encoder_hidden_states if attention_block.is_cross_attention else None,\n video_length=video_length,\n ) + hidden_states\n \n hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states\n \n output = hidden_states \n return output\n\n\nclass PositionalEncoding(nn.Module):\n def __init__(\n self, \n d_model, \n dropout = 0., \n max_len = 24\n ):\n super().__init__()\n self.dropout = nn.Dropout(p=dropout)\n position = torch.arange(max_len).unsqueeze(1)\n div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n pe = torch.zeros(1, max_len, d_model)\n pe[0, :, 0::2] = torch.sin(position * div_term)\n pe[0, :, 1::2] = torch.cos(position * div_term)\n # self.register_buffer('pe', pe)\n self.register_buffer('pe', pe, persistent=False)\n\n def forward(self, x):\n x = x + self.pe[:, :x.size(1)]\n return self.dropout(x)\n\n\nclass VersatileAttention(CrossAttention): # 继承CrossAttention类，不需要在额外写set_processor功能\n def __init__(\n self,\n attention_mode = None,\n cross_frame_attention_mode = None,\n temporal_position_encoding = False,\n temporal_position_encoding_max_len = 24, \n *args, **kwargs\n ):\n super().__init__(*args, **kwargs)\n assert attention_mode == \"Temporal\"\n\n self.attention_mode = attention_mode\n self.is_cross_attention = kwargs[\"cross_attention_dim\"] is not None\n \n self.pos_encoder = PositionalEncoding(\n kwargs[\"query_dim\"],\n dropout=0., \n max_len=temporal_position_encoding_max_len\n ) if (temporal_position_encoding and attention_mode == \"Temporal\") else None\n\n def extra_repr(self):\n return f\"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}\"\n\n def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None):\n batch_size, sequence_length, _ = hidden_states.shape\n\n if self.attention_mode == \"Temporal\":\n d = hidden_states.shape[1]\n hidden_states = rearrange(hidden_states, \"(b f) d c -> (b d) f c\", f=video_length)\n \n if self.pos_encoder is not None:\n hidden_states = self.pos_encoder(hidden_states)\n \n encoder_hidden_states = repeat(encoder_hidden_states, \"b n c -> (b d) n c\", d=d) if encoder_hidden_states is not None else encoder_hidden_states\n else:\n raise NotImplementedError\n\n encoder_hidden_states = encoder_hidden_states\n\n if self.group_norm is not None:\n hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = self.to_q(hidden_states)\n dim = query.shape[-1]\n # query = self.reshape_heads_to_batch_dim(query) # move backwards\n\n if self.added_kv_proj_dim is not None:\n raise NotImplementedError\n\n encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states\n key = self.to_k(encoder_hidden_states)\n value = self.to_v(encoder_hidden_states)\n\n ######record###### record before reshape heads to batch dim\n if self.processor is not None:\n self.processor.record_qkv(self, hidden_states, query, key, value, attention_mask)\n ##################\n\n key = self.reshape_heads_to_batch_dim(key)\n value = self.reshape_heads_to_batch_dim(value)\n\n query = self.reshape_heads_to_batch_dim(query) # reshape query here\n\n if attention_mask is not None:\n if attention_mask.shape[-1] != query.shape[1]:\n target_length = query.shape[1]\n attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)\n attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)\n\n ######record######\n # if self.processor is not None:\n # self.processor.record_attn_mask(self, hidden_states, query, key, value, attention_mask)\n ##################\n\n # attention, what we cannot get enough of\n if self._use_memory_efficient_attention_xformers:\n hidden_states = self._memory_efficient_attention_xformers(query, key, value, attention_mask)\n # Some versions of xformers return output in fp32, cast it back to the dtype of the input\n hidden_states = hidden_states.to(query.dtype)\n else:\n if self._slice_size is None or query.shape[0] // self._slice_size == 1:\n hidden_states = self._attention(query, key, value, attention_mask)\n else:\n hidden_states = self._sliced_attention(query, key, value, sequence_length, dim, attention_mask)\n\n # linear proj\n hidden_states = self.to_out[0](hidden_states)\n\n # dropout\n hidden_states = self.to_out[1](hidden_states)\n\n if self.attention_mode == \"Temporal\":\n hidden_states = rearrange(hidden_states, \"(b d) f c -> (b f) d c\", d=d)\n\n return hidden_states\n\n \n" }, { "path": "motionclone/models/resnet.py", "content": "# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/resnet.py\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom einops import rearrange\n\n\nclass InflatedConv3d(nn.Conv2d):\n def forward(self, x):\n video_length = x.shape[2]\n\n x = rearrange(x, \"b c f h w -> (b f) c h w\")\n x = super().forward(x)\n x = rearrange(x, \"(b f) c h w -> b c f h w\", f=video_length)\n\n return x\n\n\nclass InflatedGroupNorm(nn.GroupNorm):\n def forward(self, x):\n video_length = x.shape[2]\n\n x = rearrange(x, \"b c f h w -> (b f) c h w\")\n x = super().forward(x)\n x = rearrange(x, \"(b f) c h w -> b c f h w\", f=video_length)\n\n return x\n\n\nclass Upsample3D(nn.Module):\n def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name=\"conv\"):\n super().__init__()\n self.channels = channels\n self.out_channels = out_channels or channels\n self.use_conv = use_conv\n self.use_conv_transpose = use_conv_transpose\n self.name = name\n\n conv = None\n if use_conv_transpose:\n raise NotImplementedError\n elif use_conv:\n self.conv = InflatedConv3d(self.channels, self.out_channels, 3, padding=1)\n\n def forward(self, hidden_states, output_size=None):\n assert hidden_states.shape[1] == self.channels\n\n if self.use_conv_transpose:\n raise NotImplementedError\n\n # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16\n dtype = hidden_states.dtype\n if dtype == torch.bfloat16:\n hidden_states = hidden_states.to(torch.float32)\n\n # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984\n if hidden_states.shape[0] >= 64:\n hidden_states = hidden_states.contiguous()\n\n # if `output_size` is passed we force the interpolation output\n # size and do not make use of `scale_factor=2`\n if output_size is None:\n hidden_states = F.interpolate(hidden_states, scale_factor=[1.0, 2.0, 2.0], mode=\"nearest\")\n else:\n hidden_states = F.interpolate(hidden_states, size=output_size, mode=\"nearest\")\n\n # If the input is bfloat16, we cast back to bfloat16\n if dtype == torch.bfloat16:\n hidden_states = hidden_states.to(dtype)\n\n # if self.use_conv:\n # if self.name == \"conv\":\n # hidden_states = self.conv(hidden_states)\n # else:\n # hidden_states = self.Conv2d_0(hidden_states)\n hidden_states = self.conv(hidden_states)\n\n return hidden_states\n\n\nclass Downsample3D(nn.Module):\n def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name=\"conv\"):\n super().__init__()\n self.channels = channels\n self.out_channels = out_channels or channels\n self.use_conv = use_conv\n self.padding = padding\n stride = 2\n self.name = name\n\n if use_conv:\n self.conv = InflatedConv3d(self.channels, self.out_channels, 3, stride=stride, padding=padding)\n else:\n raise NotImplementedError\n\n def forward(self, hidden_states):\n assert hidden_states.shape[1] == self.channels\n if self.use_conv and self.padding == 0:\n raise NotImplementedError\n\n assert hidden_states.shape[1] == self.channels\n hidden_states = self.conv(hidden_states)\n\n return hidden_states\n\n\nclass ResnetBlock3D(nn.Module):\n def __init__(\n self,\n *,\n in_channels,\n out_channels=None,\n conv_shortcut=False,\n dropout=0.0,\n temb_channels=512,\n groups=32,\n groups_out=None,\n pre_norm=True,\n eps=1e-6,\n non_linearity=\"swish\",\n time_embedding_norm=\"default\",\n output_scale_factor=1.0,\n use_in_shortcut=None,\n use_inflated_groupnorm=False,\n ):\n super().__init__()\n self.pre_norm = pre_norm\n self.pre_norm = True\n self.in_channels = in_channels\n out_channels = in_channels if out_channels is None else out_channels\n self.out_channels = out_channels\n self.use_conv_shortcut = conv_shortcut\n self.time_embedding_norm = time_embedding_norm\n self.output_scale_factor = output_scale_factor\n self.upsample = self.downsample = None\n\n if groups_out is None:\n groups_out = groups\n\n assert use_inflated_groupnorm != None\n if use_inflated_groupnorm:\n self.norm1 = InflatedGroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)\n else:\n self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)\n\n self.conv1 = InflatedConv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)\n\n if temb_channels is not None:\n if self.time_embedding_norm == \"default\":\n time_emb_proj_out_channels = out_channels\n elif self.time_embedding_norm == \"scale_shift\":\n time_emb_proj_out_channels = out_channels * 2\n else:\n raise ValueError(f\"unknown time_embedding_norm : {self.time_embedding_norm} \")\n\n self.time_emb_proj = torch.nn.Linear(temb_channels, time_emb_proj_out_channels)\n else:\n self.time_emb_proj = None\n\n if use_inflated_groupnorm:\n self.norm2 = InflatedGroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)\n else:\n self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)\n\n self.dropout = torch.nn.Dropout(dropout)\n self.conv2 = InflatedConv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)\n\n if non_linearity == \"swish\":\n self.nonlinearity = lambda x: F.silu(x)\n elif non_linearity == \"mish\":\n self.nonlinearity = Mish()\n elif non_linearity == \"silu\":\n self.nonlinearity = nn.SiLU()\n\n self.use_in_shortcut = self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut\n\n self.conv_shortcut = None\n if self.use_in_shortcut:\n self.conv_shortcut = InflatedConv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)\n\n def forward(self, input_tensor, temb):\n hidden_states = input_tensor\n\n hidden_states = self.norm1(hidden_states)\n hidden_states = self.nonlinearity(hidden_states)\n\n hidden_states = self.conv1(hidden_states)\n\n if temb is not None:\n temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None, None]\n\n if temb is not None and self.time_embedding_norm == \"default\":\n hidden_states = hidden_states + temb\n\n hidden_states = self.norm2(hidden_states)\n\n if temb is not None and self.time_embedding_norm == \"scale_shift\":\n scale, shift = torch.chunk(temb, 2, dim=1)\n hidden_states = hidden_states * (1 + scale) + shift\n\n hidden_states = self.nonlinearity(hidden_states)\n\n hidden_states = self.dropout(hidden_states)\n hidden_states = self.conv2(hidden_states)\n\n if self.conv_shortcut is not None:\n input_tensor = self.conv_shortcut(input_tensor)\n\n output_tensor = (input_tensor + hidden_states) / self.output_scale_factor\n\n return output_tensor\n\n\nclass Mish(torch.nn.Module):\n def forward(self, hidden_states):\n return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))" }, { "path": "motionclone/models/scheduler.py", "content": "from typing import Optional, Tuple, Union\n\nimport torch\nfrom diffusers import DDIMScheduler\nfrom diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput\nfrom diffusers.utils.torch_utils import randn_tensor\n\n\nclass CustomDDIMScheduler(DDIMScheduler):\n @torch.no_grad()\n def step(\n self,\n model_output: torch.FloatTensor,\n timestep: int,\n sample: torch.FloatTensor,\n eta: float = 0.0,\n use_clipped_model_output: bool = False,\n generator=None,\n variance_noise: Optional[torch.FloatTensor] = None,\n return_dict: bool = True,\n\n # Guidance parameters\n score=None,\n guidance_scale=0.0,\n indices=None, # [0]\n\n ) -> Union[DDIMSchedulerOutput, Tuple]:\n \"\"\"\n Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion\n process from the learned model outputs (most often the predicted noise).\n\n Args:\n model_output (`torch.FloatTensor`): direct output from learned diffusion model.\n timestep (`int`): current discrete timestep in the diffusion chain.\n sample (`torch.FloatTensor`):\n current instance of sample being created by diffusion process.\n eta (`float`): weight of noise for added noise in diffusion step.\n use_clipped_model_output (`bool`): if `True`, compute \"corrected\" `model_output` from the clipped\n predicted original sample. Necessary because predicted original sample is clipped to [-1, 1] when\n `self.config.clip_sample` is `True`. If no clipping has happened, \"corrected\" `model_output` would\n coincide with the one provided as input and `use_clipped_model_output` will have not effect.\n generator: random number generator.\n variance_noise (`torch.FloatTensor`): instead of generating noise for the variance using `generator`, we\n can directly provide the noise for the variance itself. This is useful for methods such as\n CycleDiffusion. (https://arxiv.org/abs/2210.05559)\n return_dict (`bool`): option for returning tuple rather than DDIMSchedulerOutput class\n\n Returns:\n [`~schedulers.scheduling_utils.DDIMSchedulerOutput`] or `tuple`:\n [`~schedulers.scheduling_utils.DDIMSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When\n returning a tuple, the first element is the sample tensor.\n\n \"\"\"\n if self.num_inference_steps is None:\n raise ValueError(\n \"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler\"\n )\n\n # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf\n # Ideally, read DDIM paper in-detail understanding\n\n # Notation ( -> \n # - pred_noise_t -> e_theta(x_t, t)\n # - pred_original_sample -> f_theta(x_t, t) or x_0\n # - std_dev_t -> sigma_t\n # - eta -> η\n # - pred_sample_direction -> \"direction pointing to x_t\"\n # - pred_prev_sample -> \"x_t-1\"\n\n \n # Support IF models\n if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in [\"learned\", \"learned_range\"]:\n model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)\n else:\n predicted_variance = None\n\n # 1. get previous step value (=t-1)\n prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps\n\n # 2. compute alphas, betas\n alpha_prod_t = self.alphas_cumprod[timestep]\n alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod\n\n beta_prod_t = 1 - alpha_prod_t\n\n # 3. compute predicted original sample from predicted noise also called\n # \"predicted x_0\" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf\n if self.config.prediction_type == \"epsilon\":\n pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)\n pred_epsilon = model_output\n elif self.config.prediction_type == \"sample\":\n pred_original_sample = model_output\n pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)\n elif self.config.prediction_type == \"v_prediction\":\n pred_original_sample = (alpha_prod_t ** 0.5) * sample - (beta_prod_t ** 0.5) * model_output\n pred_epsilon = (alpha_prod_t ** 0.5) * model_output + (beta_prod_t ** 0.5) * sample\n else:\n raise ValueError(\n f\"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or\"\n \" `v_prediction`\"\n )\n\n # 4. Clip or threshold \"predicted x_0\"\n if self.config.thresholding:\n pred_original_sample = self._threshold_sample(pred_original_sample)\n elif self.config.clip_sample:\n pred_original_sample = pred_original_sample.clamp(\n -self.config.clip_sample_range, self.config.clip_sample_range\n )\n\n # 5. compute variance: \"sigma_t(η)\" -> see formula (16)\n # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)\n variance = self._get_variance(timestep, prev_timestep)\n std_dev_t = eta * variance ** (0.5)\n\n if use_clipped_model_output:\n # the pred_epsilon is always re-derived from the clipped x_0 in Glide\n pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) # [2, 4, 64, 64]\n\n # 6. apply guidance following the formula (14) from https://arxiv.org/pdf/2105.05233.pdf\n if score is not None and guidance_scale > 0.0: # indices指定了应用guidance的位置，此处indices = [0]\n if indices is not None:\n # import pdb; pdb.set_trace()\n assert pred_epsilon[indices].shape == score.shape, \"pred_epsilon[indices].shape != score.shape\"\n pred_epsilon[indices] = pred_epsilon[indices] - guidance_scale * (1 - alpha_prod_t) ** (0.5) * score # 只修改了其中第一个[1, 4, 64, 64]的部分\n else:\n assert pred_epsilon.shape == score.shape\n pred_epsilon = pred_epsilon - guidance_scale * (1 - alpha_prod_t) ** (0.5) * score\n # \n\n # 7. compute \"direction pointing to x_t\" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf\n pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t ** 2) ** (0.5) * pred_epsilon # [2, 4, 64, 64]\n\n # 8. compute x_t without \"random noise\" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf\n prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction # [2, 4, 64, 64]\n\n if eta > 0:\n if variance_noise is not None and generator is not None:\n raise ValueError(\n \"Cannot pass both generator and variance_noise. Please make sure that either `generator` or\"\n \" `variance_noise` stays `None`.\"\n )\n\n if variance_noise is None:\n variance_noise = randn_tensor(\n model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype\n )\n variance = std_dev_t * variance_noise # 最后还要再加一些随机噪声\n\n prev_sample = prev_sample + variance # [2, 4, 64, 64]\n self.pred_epsilon = pred_epsilon\n if not return_dict:\n return (prev_sample,)\n\n return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)\n" }, { "path": "motionclone/models/sparse_controlnet.py", "content": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# \n# Changes were made to this source code by Yuwei Guo.\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, List, Optional, Tuple, Union\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom diffusers.configuration_utils import ConfigMixin, register_to_config\nfrom diffusers.utils import BaseOutput, logging\nfrom diffusers.models.embeddings import TimestepEmbedding, Timesteps\nfrom diffusers.models.modeling_utils import ModelMixin\n\n\nfrom .unet_blocks import (\n CrossAttnDownBlock3D,\n DownBlock3D,\n UNetMidBlock3DCrossAttn,\n get_down_block,\n)\nfrom einops import repeat, rearrange\nfrom .resnet import InflatedConv3d\n\nfrom diffusers.models.unet_2d_condition import UNet2DConditionModel\n\nlogger = logging.get_logger(__name__) # pylint: disable=invalid-name\n\n\n@dataclass\nclass SparseControlNetOutput(BaseOutput):\n down_block_res_samples: Tuple[torch.Tensor]\n mid_block_res_sample: torch.Tensor\n\n\nclass SparseControlNetConditioningEmbedding(nn.Module):\n def __init__(\n self,\n conditioning_embedding_channels: int,\n conditioning_channels: int = 3,\n block_out_channels: Tuple[int] = (16, 32, 96, 256),\n ):\n super().__init__()\n\n self.conv_in = InflatedConv3d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1)\n\n self.blocks = nn.ModuleList([])\n\n for i in range(len(block_out_channels) - 1):\n channel_in = block_out_channels[i]\n channel_out = block_out_channels[i + 1]\n self.blocks.append(InflatedConv3d(channel_in, channel_in, kernel_size=3, padding=1))\n self.blocks.append(InflatedConv3d(channel_in, channel_out, kernel_size=3, padding=1, stride=2))\n\n self.conv_out = zero_module(\n InflatedConv3d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)\n )\n\n def forward(self, conditioning):\n embedding = self.conv_in(conditioning)\n embedding = F.silu(embedding)\n\n for block in self.blocks:\n embedding = block(embedding)\n embedding = F.silu(embedding)\n\n embedding = self.conv_out(embedding)\n\n return embedding\n\n\nclass SparseControlNetModel(ModelMixin, ConfigMixin):\n _supports_gradient_checkpointing = True\n\n @register_to_config\n def __init__(\n self,\n in_channels: int = 4,\n conditioning_channels: int = 3,\n flip_sin_to_cos: bool = True,\n freq_shift: int = 0,\n down_block_types: Tuple[str] = (\n \"CrossAttnDownBlock2D\",\n \"CrossAttnDownBlock2D\",\n \"CrossAttnDownBlock2D\",\n \"DownBlock2D\",\n ),\n only_cross_attention: Union[bool, Tuple[bool]] = False,\n block_out_channels: Tuple[int] = (320, 640, 1280, 1280),\n layers_per_block: int = 2,\n downsample_padding: int = 1,\n mid_block_scale_factor: float = 1,\n act_fn: str = \"silu\",\n norm_num_groups: Optional[int] = 32,\n norm_eps: float = 1e-5,\n cross_attention_dim: int = 1280,\n attention_head_dim: Union[int, Tuple[int]] = 8,\n num_attention_heads: Optional[Union[int, Tuple[int]]] = None,\n use_linear_projection: bool = False,\n class_embed_type: Optional[str] = None,\n num_class_embeds: Optional[int] = None,\n upcast_attention: bool = False,\n resnet_time_scale_shift: str = \"default\",\n projection_class_embeddings_input_dim: Optional[int] = None,\n controlnet_conditioning_channel_order: str = \"rgb\",\n conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),\n global_pool_conditions: bool = False,\n\n use_motion_module = True,\n motion_module_resolutions = ( 1,2,4,8 ),\n motion_module_mid_block = False,\n motion_module_type = \"Vanilla\",\n motion_module_kwargs = {\n \"num_attention_heads\": 8,\n \"num_transformer_block\": 1,\n \"attention_block_types\": [\"Temporal_Self\"],\n \"temporal_position_encoding\": True,\n \"temporal_position_encoding_max_len\": 32,\n \"temporal_attention_dim_div\": 1,\n \"causal_temporal_attention\": False,\n },\n\n concate_conditioning_mask: bool = True,\n use_simplified_condition_embedding: bool = False,\n\n set_noisy_sample_input_to_zero: bool = False,\n ):\n super().__init__()\n\n # If `num_attention_heads` is not defined (which is the case for most models)\n # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.\n # The reason for this behavior is to correct for incorrectly named variables that were introduced\n # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131\n # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking\n # which is why we correct for the naming here.\n num_attention_heads = num_attention_heads or attention_head_dim\n\n # Check inputs\n if len(block_out_channels) != len(down_block_types):\n raise ValueError(\n 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}.\"\n )\n\n if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):\n raise ValueError(\n f\"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}.\"\n )\n\n if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):\n raise ValueError(\n 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}.\"\n )\n\n # input\n self.set_noisy_sample_input_to_zero = set_noisy_sample_input_to_zero\n\n conv_in_kernel = 3\n conv_in_padding = (conv_in_kernel - 1) // 2\n self.conv_in = InflatedConv3d(\n in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding\n )\n\n if concate_conditioning_mask:\n conditioning_channels = conditioning_channels + 1\n self.concate_conditioning_mask = concate_conditioning_mask\n\n # control net conditioning embedding\n if use_simplified_condition_embedding:\n self.controlnet_cond_embedding = zero_module(\n InflatedConv3d(conditioning_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding)\n ).to(torch.float16)\n else:\n self.controlnet_cond_embedding = SparseControlNetConditioningEmbedding(\n conditioning_embedding_channels=block_out_channels[0],\n block_out_channels=conditioning_embedding_out_channels,\n conditioning_channels=conditioning_channels,\n ).to(torch.float16)\n self.use_simplified_condition_embedding = use_simplified_condition_embedding\n\n # time\n time_embed_dim = block_out_channels[0] * 4\n\n self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)\n timestep_input_dim = block_out_channels[0]\n\n self.time_embedding = TimestepEmbedding(\n timestep_input_dim,\n time_embed_dim,\n act_fn=act_fn,\n )\n\n # class embedding\n if class_embed_type is None and num_class_embeds is not None:\n self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)\n elif class_embed_type == \"timestep\":\n self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)\n elif class_embed_type == \"identity\":\n self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)\n elif class_embed_type == \"projection\":\n if projection_class_embeddings_input_dim is None:\n raise ValueError(\n \"`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set\"\n )\n # The projection `class_embed_type` is the same as the timestep `class_embed_type` except\n # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings\n # 2. it projects from an arbitrary input dimension.\n #\n # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.\n # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.\n # As a result, `TimestepEmbedding` can be passed arbitrary vectors.\n self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)\n else:\n self.class_embedding = None\n\n\n self.down_blocks = nn.ModuleList([])\n self.controlnet_down_blocks = nn.ModuleList([])\n\n if isinstance(only_cross_attention, bool):\n only_cross_attention = [only_cross_attention] * len(down_block_types)\n\n if isinstance(attention_head_dim, int):\n attention_head_dim = (attention_head_dim,) * len(down_block_types)\n\n if isinstance(num_attention_heads, int):\n num_attention_heads = (num_attention_heads,) * len(down_block_types)\n\n # down\n output_channel = block_out_channels[0]\n\n controlnet_block = InflatedConv3d(output_channel, output_channel, kernel_size=1)\n controlnet_block = zero_module(controlnet_block)\n self.controlnet_down_blocks.append(controlnet_block)\n\n for i, down_block_type in enumerate(down_block_types):\n res = 2 ** i\n input_channel = output_channel\n output_channel = block_out_channels[i]\n is_final_block = i == len(block_out_channels) - 1\n\n down_block = get_down_block(\n down_block_type,\n num_layers=layers_per_block,\n in_channels=input_channel,\n out_channels=output_channel,\n temb_channels=time_embed_dim,\n add_downsample=not is_final_block,\n resnet_eps=norm_eps,\n resnet_act_fn=act_fn,\n resnet_groups=norm_num_groups,\n cross_attention_dim=cross_attention_dim,\n attn_num_head_channels=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,\n downsample_padding=downsample_padding,\n use_linear_projection=use_linear_projection,\n only_cross_attention=only_cross_attention[i],\n upcast_attention=upcast_attention,\n resnet_time_scale_shift=resnet_time_scale_shift,\n\n use_inflated_groupnorm=True,\n\n use_motion_module=use_motion_module and (res in motion_module_resolutions),\n motion_module_type=motion_module_type,\n motion_module_kwargs=motion_module_kwargs,\n )\n self.down_blocks.append(down_block)\n\n for _ in range(layers_per_block):\n controlnet_block = InflatedConv3d(output_channel, output_channel, kernel_size=1)\n controlnet_block = zero_module(controlnet_block)\n self.controlnet_down_blocks.append(controlnet_block)\n\n if not is_final_block:\n controlnet_block = InflatedConv3d(output_channel, output_channel, kernel_size=1)\n controlnet_block = zero_module(controlnet_block)\n self.controlnet_down_blocks.append(controlnet_block)\n\n # mid\n mid_block_channel = block_out_channels[-1]\n\n controlnet_block = InflatedConv3d(mid_block_channel, mid_block_channel, kernel_size=1)\n controlnet_block = zero_module(controlnet_block)\n self.controlnet_mid_block = controlnet_block\n\n self.mid_block = UNetMidBlock3DCrossAttn(\n in_channels=mid_block_channel,\n temb_channels=time_embed_dim,\n resnet_eps=norm_eps,\n resnet_act_fn=act_fn,\n output_scale_factor=mid_block_scale_factor,\n resnet_time_scale_shift=resnet_time_scale_shift,\n cross_attention_dim=cross_attention_dim,\n attn_num_head_channels=num_attention_heads[-1],\n resnet_groups=norm_num_groups,\n use_linear_projection=use_linear_projection,\n upcast_attention=upcast_attention,\n\n use_inflated_groupnorm=True,\n use_motion_module=use_motion_module and motion_module_mid_block,\n motion_module_type=motion_module_type,\n motion_module_kwargs=motion_module_kwargs,\n )\n\n @classmethod\n def from_unet(\n cls,\n unet: UNet2DConditionModel,\n controlnet_conditioning_channel_order: str = \"rgb\",\n conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),\n load_weights_from_unet: bool = True,\n\n controlnet_additional_kwargs: dict = {},\n ):\n controlnet = cls(\n in_channels=unet.config.in_channels,\n flip_sin_to_cos=unet.config.flip_sin_to_cos,\n freq_shift=unet.config.freq_shift,\n down_block_types=unet.config.down_block_types,\n only_cross_attention=unet.config.only_cross_attention,\n block_out_channels=unet.config.block_out_channels,\n layers_per_block=unet.config.layers_per_block,\n downsample_padding=unet.config.downsample_padding,\n mid_block_scale_factor=unet.config.mid_block_scale_factor,\n act_fn=unet.config.act_fn,\n norm_num_groups=unet.config.norm_num_groups,\n norm_eps=unet.config.norm_eps,\n cross_attention_dim=unet.config.cross_attention_dim,\n attention_head_dim=unet.config.attention_head_dim,\n num_attention_heads=unet.config.num_attention_heads,\n use_linear_projection=unet.config.use_linear_projection,\n class_embed_type=unet.config.class_embed_type,\n num_class_embeds=unet.config.num_class_embeds,\n upcast_attention=unet.config.upcast_attention,\n resnet_time_scale_shift=unet.config.resnet_time_scale_shift,\n projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,\n controlnet_conditioning_channel_order=controlnet_conditioning_channel_order,\n conditioning_embedding_out_channels=conditioning_embedding_out_channels,\n\n **controlnet_additional_kwargs,\n )\n\n if load_weights_from_unet:\n m, u = controlnet.conv_in.load_state_dict(cls.image_layer_filter(unet.conv_in.state_dict()), strict=False)\n assert len(u) == 0\n m, u = controlnet.time_proj.load_state_dict(cls.image_layer_filter(unet.time_proj.state_dict()), strict=False)\n assert len(u) == 0\n m, u = controlnet.time_embedding.load_state_dict(cls.image_layer_filter(unet.time_embedding.state_dict()), strict=False)\n assert len(u) == 0\n\n if controlnet.class_embedding:\n m, u = controlnet.class_embedding.load_state_dict(cls.image_layer_filter(unet.class_embedding.state_dict()), strict=False)\n assert len(u) == 0\n m, u = controlnet.down_blocks.load_state_dict(cls.image_layer_filter(unet.down_blocks.state_dict()), strict=False)\n assert len(u) == 0\n m, u = controlnet.mid_block.load_state_dict(cls.image_layer_filter(unet.mid_block.state_dict()), strict=False)\n assert len(u) == 0\n\n return controlnet\n\n @staticmethod\n def image_layer_filter(state_dict):\n new_state_dict = {}\n for name, param in state_dict.items():\n if \"motion_modules.\" in name or \"lora\" in name: continue\n new_state_dict[name] = param\n return new_state_dict\n\n # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice\n def set_attention_slice(self, slice_size):\n r\"\"\"\n Enable sliced attention computation.\n\n When this option is enabled, the attention module splits the input tensor in slices to compute attention in\n several steps. This is useful for saving some memory in exchange for a small decrease in speed.\n\n Args:\n slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `\"auto\"`):\n When `\"auto\"`, input to the attention heads is halved, so attention is computed in two steps. If\n `\"max\"`, maximum amount of memory is saved by running only one slice at a time. If a number is\n provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`\n must be a multiple of `slice_size`.\n \"\"\"\n sliceable_head_dims = []\n\n def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):\n if hasattr(module, \"set_attention_slice\"):\n sliceable_head_dims.append(module.sliceable_head_dim)\n\n for child in module.children():\n fn_recursive_retrieve_sliceable_dims(child)\n\n # retrieve number of attention layers\n for module in self.children():\n fn_recursive_retrieve_sliceable_dims(module)\n\n num_sliceable_layers = len(sliceable_head_dims)\n\n if slice_size == \"auto\":\n # half the attention head size is usually a good trade-off between\n # speed and memory\n slice_size = [dim // 2 for dim in sliceable_head_dims]\n elif slice_size == \"max\":\n # make smallest slice possible\n slice_size = num_sliceable_layers * [1]\n\n slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size\n\n if len(slice_size) != len(sliceable_head_dims):\n raise ValueError(\n f\"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different\"\n f\" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}.\"\n )\n\n for i in range(len(slice_size)):\n size = slice_size[i]\n dim = sliceable_head_dims[i]\n if size is not None and size > dim:\n raise ValueError(f\"size {size} has to be smaller or equal to {dim}.\")\n\n # Recursively walk through all the children.\n # Any children which exposes the set_attention_slice method\n # gets the message\n def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):\n if hasattr(module, \"set_attention_slice\"):\n module.set_attention_slice(slice_size.pop())\n\n for child in module.children():\n fn_recursive_set_attention_slice(child, slice_size)\n\n reversed_slice_size = list(reversed(slice_size))\n for module in self.children():\n fn_recursive_set_attention_slice(module, reversed_slice_size)\n\n def _set_gradient_checkpointing(self, module, value=False):\n if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)):\n module.gradient_checkpointing = value\n\n def forward(\n self,\n sample: torch.FloatTensor,\n timestep: Union[torch.Tensor, float, int],\n encoder_hidden_states: torch.Tensor,\n\n controlnet_cond: torch.FloatTensor,\n conditioning_mask: Optional[torch.FloatTensor] = None,\n\n conditioning_scale: float = 1.0,\n class_labels: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n cross_attention_kwargs: Optional[Dict[str, Any]] = None,\n guess_mode: bool = False,\n return_dict: bool = True,\n ) -> Union[SparseControlNetOutput, Tuple]:\n\n # set input noise to zero\n # if self.set_noisy_sample_input_to_zero:\n # sample = torch.zeros_like(sample).to(sample.device)\n\n # prepare attention_mask\n if attention_mask is not None:\n attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0\n attention_mask = attention_mask.unsqueeze(1)\n\n # 1. time\n timesteps = timestep\n if not torch.is_tensor(timesteps):\n # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can\n # This would be a good case for the `match` statement (Python 3.10+)\n is_mps = sample.device.type == \"mps\"\n if isinstance(timestep, float):\n dtype = torch.float32 if is_mps else torch.float64\n else:\n dtype = torch.int32 if is_mps else torch.int64\n timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)\n elif len(timesteps.shape) == 0:\n timesteps = timesteps[None].to(sample.device)\n\n timesteps = timesteps.repeat(sample.shape[0] // timesteps.shape[0])\n encoder_hidden_states = encoder_hidden_states.repeat(sample.shape[0] // encoder_hidden_states.shape[0], 1, 1)\n\n # broadcast to batch dimension in a way that's compatible with ONNX/Core ML\n timesteps = timesteps.expand(sample.shape[0])\n\n t_emb = self.time_proj(timesteps)\n\n # timesteps does not contain any weights and will always return f32 tensors\n # but time_embedding might actually be running in fp16. so we need to cast here.\n # there might be better ways to encapsulate this.\n t_emb = t_emb.to(dtype=self.dtype)\n emb = self.time_embedding(t_emb)\n\n if self.class_embedding is not None:\n if class_labels is None:\n raise ValueError(\"class_labels should be provided when num_class_embeds > 0\")\n\n if self.config.class_embed_type == \"timestep\":\n class_labels = self.time_proj(class_labels)\n\n class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)\n emb = emb + class_emb\n\n # 2. pre-process\n # equal to set input noise to zero\n if self.set_noisy_sample_input_to_zero:\n shape = sample.shape\n sample = self.conv_in.bias.reshape(1,-1,1,1,1).expand(shape[0],-1,shape[2],shape[3],shape[4])\n else:\n sample = self.conv_in(sample)\n\n if self.concate_conditioning_mask:\n controlnet_cond = torch.cat([controlnet_cond, conditioning_mask], dim=1).to(torch.float16)\n # import pdb; pdb.set_trace()\n controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)\n \n sample = sample + controlnet_cond\n\n # 3. down\n down_block_res_samples = (sample,)\n for downsample_block in self.down_blocks:\n if hasattr(downsample_block, \"has_cross_attention\") and downsample_block.has_cross_attention:\n sample, res_samples = downsample_block(\n hidden_states=sample,\n temb=emb,\n encoder_hidden_states=encoder_hidden_states,\n attention_mask=attention_mask,\n # cross_attention_kwargs=cross_attention_kwargs,\n )\n else: sample, res_samples = downsample_block(hidden_states=sample, temb=emb)\n\n down_block_res_samples += res_samples\n\n # 4. mid\n if self.mid_block is not None:\n sample = self.mid_block(\n sample,\n emb,\n encoder_hidden_states=encoder_hidden_states,\n attention_mask=attention_mask,\n # cross_attention_kwargs=cross_attention_kwargs,\n )\n\n # 5. controlnet blocks\n controlnet_down_block_res_samples = ()\n\n for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):\n down_block_res_sample = controlnet_block(down_block_res_sample)\n controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,)\n\n down_block_res_samples = controlnet_down_block_res_samples\n\n mid_block_res_sample = self.controlnet_mid_block(sample)\n\n # 6. scaling\n if guess_mode and not self.config.global_pool_conditions:\n scales = torch.logspace(-1, 0, len(down_block_res_samples) + 1, device=sample.device) # 0.1 to 1.0\n\n scales = scales * conditioning_scale\n down_block_res_samples = [sample * scale for sample, scale in zip(down_block_res_samples, scales)]\n mid_block_res_sample = mid_block_res_sample * scales[-1] # last one\n else:\n down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples]\n mid_block_res_sample = mid_block_res_sample * conditioning_scale\n\n if self.config.global_pool_conditions:\n down_block_res_samples = [\n torch.mean(sample, dim=(2, 3), keepdim=True) for sample in down_block_res_samples\n ]\n mid_block_res_sample = torch.mean(mid_block_res_sample, dim=(2, 3), keepdim=True)\n\n if not return_dict:\n return (down_block_res_samples, mid_block_res_sample)\n\n return SparseControlNetOutput(\n down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample\n )\n\n\ndef zero_module(module):\n for p in module.parameters():\n nn.init.zeros_(p)\n return module\n" }, { "path": "motionclone/models/unet.py", "content": "# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_condition.py\n\nfrom dataclasses import dataclass\nfrom typing import List, Optional, Tuple, Union\n\nimport os\nimport json\nimport pdb\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.checkpoint\n\nfrom diffusers.configuration_utils import ConfigMixin, register_to_config\nfrom diffusers.models.modeling_utils import ModelMixin\nfrom diffusers.utils import BaseOutput, logging\nfrom diffusers.models.embeddings import TimestepEmbedding, Timesteps\nfrom .unet_blocks import (\n CrossAttnDownBlock3D,\n CrossAttnUpBlock3D,\n DownBlock3D,\n UNetMidBlock3DCrossAttn,\n UpBlock3D,\n get_down_block,\n get_up_block,\n)\nfrom .resnet import InflatedConv3d, InflatedGroupNorm\n\n\nlogger = logging.get_logger(__name__) # pylint: disable=invalid-name\n\n\n@dataclass\nclass UNet3DConditionOutput(BaseOutput):\n sample: torch.FloatTensor\n\n\nclass UNet3DConditionModel(ModelMixin, ConfigMixin):\n _supports_gradient_checkpointing = True\n\n @register_to_config\n def __init__(\n self,\n sample_size: Optional[int] = None,\n in_channels: int = 4,\n out_channels: int = 4,\n center_input_sample: bool = False,\n flip_sin_to_cos: bool = True,\n freq_shift: int = 0, \n down_block_types: Tuple[str] = (\n \"CrossAttnDownBlock3D\",\n \"CrossAttnDownBlock3D\",\n \"CrossAttnDownBlock3D\",\n \"DownBlock3D\",\n ),\n mid_block_type: str = \"UNetMidBlock3DCrossAttn\",\n up_block_types: Tuple[str] = ( # 第一个不带有CrossAttn，后面三个带有CrossAttn\n \"UpBlock3D\",\n \"CrossAttnUpBlock3D\",\n \"CrossAttnUpBlock3D\",\n \"CrossAttnUpBlock3D\"\n ),\n only_cross_attention: Union[bool, Tuple[bool]] = False,\n block_out_channels: Tuple[int] = (320, 640, 1280, 1280),\n layers_per_block: int = 2,\n downsample_padding: int = 1,\n mid_block_scale_factor: float = 1,\n act_fn: str = \"silu\",\n norm_num_groups: int = 32,\n norm_eps: float = 1e-5,\n cross_attention_dim: int = 1280,\n attention_head_dim: Union[int, Tuple[int]] = 8,\n dual_cross_attention: bool = False,\n use_linear_projection: bool = False,\n class_embed_type: Optional[str] = None,\n num_class_embeds: Optional[int] = None,\n upcast_attention: bool = False,\n resnet_time_scale_shift: str = \"default\",\n \n use_inflated_groupnorm=False,\n \n # Additional\n use_motion_module = False,\n motion_module_resolutions = ( 1,2,4,8 ),\n motion_module_mid_block = False,\n motion_module_decoder_only = False,\n motion_module_type = None,\n motion_module_kwargs = {},\n unet_use_cross_frame_attention = False,\n unet_use_temporal_attention = False,\n ):\n super().__init__()\n \n self.sample_size = sample_size\n time_embed_dim = block_out_channels[0] * 4\n\n # input\n self.conv_in = InflatedConv3d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1))\n\n # time\n self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)\n timestep_input_dim = block_out_channels[0]\n\n self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)\n\n # class embedding\n if class_embed_type is None and num_class_embeds is not None:\n self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)\n elif class_embed_type == \"timestep\":\n self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)\n elif class_embed_type == \"identity\":\n self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)\n else:\n self.class_embedding = None\n\n self.down_blocks = nn.ModuleList([])\n self.mid_block = None\n self.up_blocks = nn.ModuleList([])\n\n if isinstance(only_cross_attention, bool):\n only_cross_attention = [only_cross_attention] * len(down_block_types)\n\n if isinstance(attention_head_dim, int):\n attention_head_dim = (attention_head_dim,) * len(down_block_types)\n\n # down\n output_channel = block_out_channels[0]\n for i, down_block_type in enumerate(down_block_types):\n res = 2 ** i\n input_channel = output_channel\n output_channel = block_out_channels[i]\n is_final_block = i == len(block_out_channels) - 1\n\n down_block = get_down_block(\n down_block_type,\n num_layers=layers_per_block,\n in_channels=input_channel,\n out_channels=output_channel,\n temb_channels=time_embed_dim,\n add_downsample=not is_final_block,\n resnet_eps=norm_eps,\n resnet_act_fn=act_fn,\n resnet_groups=norm_num_groups,\n cross_attention_dim=cross_attention_dim,\n attn_num_head_channels=attention_head_dim[i],\n downsample_padding=downsample_padding,\n dual_cross_attention=dual_cross_attention,\n use_linear_projection=use_linear_projection,\n only_cross_attention=only_cross_attention[i],\n upcast_attention=upcast_attention,\n resnet_time_scale_shift=resnet_time_scale_shift,\n\n unet_use_cross_frame_attention=unet_use_cross_frame_attention,\n unet_use_temporal_attention=unet_use_temporal_attention,\n use_inflated_groupnorm=use_inflated_groupnorm,\n \n use_motion_module=use_motion_module and (res in motion_module_resolutions) and (not motion_module_decoder_only),\n motion_module_type=motion_module_type,\n motion_module_kwargs=motion_module_kwargs,\n )\n self.down_blocks.append(down_block)\n\n # mid\n if mid_block_type == \"UNetMidBlock3DCrossAttn\":\n self.mid_block = UNetMidBlock3DCrossAttn(\n in_channels=block_out_channels[-1],\n temb_channels=time_embed_dim,\n resnet_eps=norm_eps,\n resnet_act_fn=act_fn,\n output_scale_factor=mid_block_scale_factor,\n resnet_time_scale_shift=resnet_time_scale_shift,\n cross_attention_dim=cross_attention_dim,\n attn_num_head_channels=attention_head_dim[-1],\n resnet_groups=norm_num_groups,\n dual_cross_attention=dual_cross_attention,\n use_linear_projection=use_linear_projection,\n upcast_attention=upcast_attention,\n\n unet_use_cross_frame_attention=unet_use_cross_frame_attention,\n unet_use_temporal_attention=unet_use_temporal_attention,\n use_inflated_groupnorm=use_inflated_groupnorm,\n \n use_motion_module=use_motion_module and motion_module_mid_block,\n motion_module_type=motion_module_type,\n motion_module_kwargs=motion_module_kwargs,\n )\n else:\n raise ValueError(f\"unknown mid_block_type : {mid_block_type}\")\n \n # count how many layers upsample the videos\n self.num_upsamplers = 0\n\n # up\n reversed_block_out_channels = list(reversed(block_out_channels))\n reversed_attention_head_dim = list(reversed(attention_head_dim))\n only_cross_attention = list(reversed(only_cross_attention))\n output_channel = reversed_block_out_channels[0]\n for i, up_block_type in enumerate(up_block_types):\n res = 2 ** (3 - i)\n is_final_block = i == len(block_out_channels) - 1\n\n prev_output_channel = output_channel\n output_channel = reversed_block_out_channels[i]\n input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]\n\n # add upsample block for all BUT final layer\n if not is_final_block:\n add_upsample = True\n self.num_upsamplers += 1\n else:\n add_upsample = False\n\n up_block = get_up_block(\n up_block_type,\n num_layers=layers_per_block + 1,\n in_channels=input_channel,\n out_channels=output_channel,\n prev_output_channel=prev_output_channel,\n temb_channels=time_embed_dim,\n add_upsample=add_upsample,\n resnet_eps=norm_eps,\n resnet_act_fn=act_fn,\n resnet_groups=norm_num_groups,\n cross_attention_dim=cross_attention_dim,\n attn_num_head_channels=reversed_attention_head_dim[i],\n dual_cross_attention=dual_cross_attention,\n use_linear_projection=use_linear_projection,\n only_cross_attention=only_cross_attention[i],\n upcast_attention=upcast_attention,\n resnet_time_scale_shift=resnet_time_scale_shift,\n\n unet_use_cross_frame_attention=unet_use_cross_frame_attention,\n unet_use_temporal_attention=unet_use_temporal_attention,\n use_inflated_groupnorm=use_inflated_groupnorm,\n\n use_motion_module=use_motion_module and (res in motion_module_resolutions),\n motion_module_type=motion_module_type,\n motion_module_kwargs=motion_module_kwargs,\n )\n self.up_blocks.append(up_block)\n prev_output_channel = output_channel\n\n # out\n if use_inflated_groupnorm:\n self.conv_norm_out = InflatedGroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps)\n else:\n self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps)\n self.conv_act = nn.SiLU()\n self.conv_out = InflatedConv3d(block_out_channels[0], out_channels, kernel_size=3, padding=1)\n\n def set_attention_slice(self, slice_size):\n r\"\"\"\n Enable sliced attention computation.\n\n When this option is enabled, the attention module will split the input tensor in slices, to compute attention\n in several steps. This is useful to save some memory in exchange for a small speed decrease.\n\n Args:\n slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `\"auto\"`):\n When `\"auto\"`, halves the input to the attention heads, so attention will be computed in two steps. If\n `\"max\"`, maxium amount of memory will be saved by running only one slice at a time. If a number is\n provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`\n must be a multiple of `slice_size`.\n \"\"\"\n sliceable_head_dims = []\n\n def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):\n if hasattr(module, \"set_attention_slice\"):\n sliceable_head_dims.append(module.sliceable_head_dim)\n\n for child in module.children():\n fn_recursive_retrieve_slicable_dims(child)\n\n # retrieve number of attention layers\n for module in self.children():\n fn_recursive_retrieve_slicable_dims(module)\n\n num_slicable_layers = len(sliceable_head_dims)\n\n if slice_size == \"auto\":\n # half the attention head size is usually a good trade-off between\n # speed and memory\n slice_size = [dim // 2 for dim in sliceable_head_dims]\n elif slice_size == \"max\":\n # make smallest slice possible\n slice_size = num_slicable_layers * [1]\n\n slice_size = num_slicable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size\n\n if len(slice_size) != len(sliceable_head_dims):\n raise ValueError(\n f\"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different\"\n f\" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}.\"\n )\n\n for i in range(len(slice_size)):\n size = slice_size[i]\n dim = sliceable_head_dims[i]\n if size is not None and size > dim:\n raise ValueError(f\"size {size} has to be smaller or equal to {dim}.\")\n\n # Recursively walk through all the children.\n # Any children which exposes the set_attention_slice method\n # gets the message\n def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):\n if hasattr(module, \"set_attention_slice\"):\n module.set_attention_slice(slice_size.pop())\n\n for child in module.children():\n fn_recursive_set_attention_slice(child, slice_size)\n\n reversed_slice_size = list(reversed(slice_size))\n for module in self.children():\n fn_recursive_set_attention_slice(module, reversed_slice_size)\n\n def _set_gradient_checkpointing(self, module, value=False):\n if isinstance(module, (CrossAttnDownBlock3D, DownBlock3D, CrossAttnUpBlock3D, UpBlock3D)):\n module.gradient_checkpointing = value\n\n def forward(\n self,\n sample: torch.FloatTensor,\n timestep: Union[torch.Tensor, float, int],\n encoder_hidden_states: torch.Tensor,\n class_labels: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n\n # support controlnet\n down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,\n mid_block_additional_residual: Optional[torch.Tensor] = None,\n\n return_dict: bool = True,\n ) -> Union[UNet3DConditionOutput, Tuple]:\n r\"\"\"\n Args:\n sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor\n timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps\n encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states\n return_dict (`bool`, *optional*, defaults to `True`):\n Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.\n\n Returns:\n [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:\n [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When\n returning a tuple, the first element is the sample tensor.\n \"\"\"\n # By default samples have to be AT least a multiple of the overall upsampling factor.\n # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).\n # However, the upsampling interpolation output size can be forced to fit any upsampling size\n # on the fly if necessary.\n default_overall_up_factor = 2**self.num_upsamplers\n\n # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`\n forward_upsample_size = False\n upsample_size = None\n\n if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):\n logger.info(\"Forward upsample size to force interpolation output size.\")\n forward_upsample_size = True\n\n # prepare attention_mask\n if attention_mask is not None:\n attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0\n attention_mask = attention_mask.unsqueeze(1)\n\n # center input if necessary\n if self.config.center_input_sample:\n sample = 2 * sample - 1.0\n\n # time\n timesteps = timestep\n if not torch.is_tensor(timesteps):\n # This would be a good case for the `match` statement (Python 3.10+)\n is_mps = sample.device.type == \"mps\"\n if isinstance(timestep, float):\n dtype = torch.float32 if is_mps else torch.float64\n else:\n dtype = torch.int32 if is_mps else torch.int64\n timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)\n elif len(timesteps.shape) == 0:\n timesteps = timesteps[None].to(sample.device)\n\n # broadcast to batch dimension in a way that's compatible with ONNX/Core ML\n timesteps = timesteps.expand(sample.shape[0])\n\n t_emb = self.time_proj(timesteps)\n\n # timesteps does not contain any weights and will always return f32 tensors\n # but time_embedding might actually be running in fp16. so we need to cast here.\n # there might be better ways to encapsulate this.\n t_emb = t_emb.to(dtype=self.dtype)\n emb = self.time_embedding(t_emb)\n\n if self.class_embedding is not None:\n if class_labels is None:\n raise ValueError(\"class_labels should be provided when num_class_embeds > 0\")\n\n if self.config.class_embed_type == \"timestep\":\n class_labels = self.time_proj(class_labels)\n\n class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)\n emb = emb + class_emb\n\n # pre-process\n sample = self.conv_in(sample)\n\n # down\n down_block_res_samples = (sample,)\n for downsample_block in self.down_blocks:\n if hasattr(downsample_block, \"has_cross_attention\") and downsample_block.has_cross_attention:\n sample, res_samples = downsample_block(\n hidden_states=sample,\n temb=emb,\n encoder_hidden_states=encoder_hidden_states,\n attention_mask=attention_mask,\n )\n else:\n sample, res_samples = downsample_block(hidden_states=sample, temb=emb, encoder_hidden_states=encoder_hidden_states)\n\n down_block_res_samples += res_samples\n\n # support controlnet\n down_block_res_samples = list(down_block_res_samples)\n if down_block_additional_residuals is not None:\n for i, down_block_additional_residual in enumerate(down_block_additional_residuals):\n if down_block_additional_residual.dim() == 4: # boardcast\n down_block_additional_residual = down_block_additional_residual.unsqueeze(2)\n down_block_res_samples[i] = down_block_res_samples[i] + down_block_additional_residual\n\n # mid\n sample = self.mid_block(\n sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask\n )\n\n # support controlnet\n if mid_block_additional_residual is not None:\n if mid_block_additional_residual.dim() == 4: # boardcast\n mid_block_additional_residual = mid_block_additional_residual.unsqueeze(2)\n sample = sample + mid_block_additional_residual\n\n # up\n for i, upsample_block in enumerate(self.up_blocks):\n is_final_block = i == len(self.up_blocks) - 1\n\n res_samples = down_block_res_samples[-len(upsample_block.resnets) :]\n down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]\n\n # if we have not reached the final block and need to forward the\n # upsample size, we do it here\n if not is_final_block and forward_upsample_size:\n upsample_size = down_block_res_samples[-1].shape[2:]\n\n if hasattr(upsample_block, \"has_cross_attention\") and upsample_block.has_cross_attention:\n sample = upsample_block(\n hidden_states=sample,\n temb=emb,\n res_hidden_states_tuple=res_samples,\n encoder_hidden_states=encoder_hidden_states,\n upsample_size=upsample_size,\n attention_mask=attention_mask,\n )\n else:\n sample = upsample_block(\n hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size, encoder_hidden_states=encoder_hidden_states,\n )\n\n # post-process\n sample = self.conv_norm_out(sample)\n sample = self.conv_act(sample)\n sample = self.conv_out(sample)\n\n if not return_dict:\n return (sample,)\n\n return UNet3DConditionOutput(sample=sample)\n\n @classmethod\n def from_pretrained_2d(cls, pretrained_model_path, subfolder=None, unet_additional_kwargs=None):\n if subfolder is not None:\n pretrained_model_path = os.path.join(pretrained_model_path, subfolder)\n print(f\"loaded 3D unet's pretrained weights from {pretrained_model_path} ...\")\n\n config_file = os.path.join(pretrained_model_path, 'config.json')\n if not os.path.isfile(config_file):\n raise RuntimeError(f\"{config_file} does not exist\")\n with open(config_file, \"r\") as f:\n config = json.load(f)\n config[\"_class_name\"] = cls.__name__\n config[\"down_block_types\"] = [\n \"CrossAttnDownBlock3D\",\n \"CrossAttnDownBlock3D\",\n \"CrossAttnDownBlock3D\",\n \"DownBlock3D\"\n ]\n config[\"up_block_types\"] = [\n \"UpBlock3D\",\n \"CrossAttnUpBlock3D\",\n \"CrossAttnUpBlock3D\",\n \"CrossAttnUpBlock3D\"\n ]\n\n from diffusers.utils import WEIGHTS_NAME\n model = cls.from_config(config, **unet_additional_kwargs)\n model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)\n if not os.path.isfile(model_file):\n raise RuntimeError(f\"{model_file} does not exist\")\n state_dict = torch.load(model_file, map_location=\"cpu\")\n\n m, u = model.load_state_dict(state_dict, strict=False)\n print(f\"### motion keys will be loaded: {len(m)}; \\n### unexpected keys: {len(u)};\")\n \n params = [p.numel() if \"motion_modules.\" in n else 0 for n, p in model.named_parameters()]\n print(f\"### Motion Module Parameters: {sum(params) / 1e6} M\")\n \n return model\n" }, { "path": "motionclone/models/unet_blocks.py", "content": "# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_blocks.py\n\nimport torch\nfrom torch import nn\n\nfrom .attention import Transformer3DModel\nfrom .resnet import Downsample3D, ResnetBlock3D, Upsample3D\nfrom .motion_module import get_motion_module\n\nimport pdb\n\ndef get_down_block(\n down_block_type,\n num_layers,\n in_channels,\n out_channels,\n temb_channels,\n add_downsample,\n resnet_eps,\n resnet_act_fn,\n attn_num_head_channels,\n resnet_groups=None,\n cross_attention_dim=None,\n downsample_padding=None,\n dual_cross_attention=False,\n use_linear_projection=False,\n only_cross_attention=False,\n upcast_attention=False,\n resnet_time_scale_shift=\"default\",\n \n unet_use_cross_frame_attention=False,\n unet_use_temporal_attention=False,\n use_inflated_groupnorm=False,\n\n use_motion_module=None,\n \n motion_module_type=None,\n motion_module_kwargs=None,\n):\n down_block_type = down_block_type[7:] if down_block_type.startswith(\"UNetRes\") else down_block_type\n if down_block_type == \"DownBlock3D\":\n return DownBlock3D(\n num_layers=num_layers,\n in_channels=in_channels,\n out_channels=out_channels,\n temb_channels=temb_channels,\n add_downsample=add_downsample,\n resnet_eps=resnet_eps,\n resnet_act_fn=resnet_act_fn,\n resnet_groups=resnet_groups,\n downsample_padding=downsample_padding,\n resnet_time_scale_shift=resnet_time_scale_shift,\n\n use_inflated_groupnorm=use_inflated_groupnorm,\n\n use_motion_module=use_motion_module,\n motion_module_type=motion_module_type,\n motion_module_kwargs=motion_module_kwargs,\n )\n elif down_block_type == \"CrossAttnDownBlock3D\":\n if cross_attention_dim is None:\n raise ValueError(\"cross_attention_dim must be specified for CrossAttnDownBlock3D\")\n return CrossAttnDownBlock3D(\n num_layers=num_layers,\n in_channels=in_channels,\n out_channels=out_channels,\n temb_channels=temb_channels,\n add_downsample=add_downsample,\n resnet_eps=resnet_eps,\n resnet_act_fn=resnet_act_fn,\n resnet_groups=resnet_groups,\n downsample_padding=downsample_padding,\n cross_attention_dim=cross_attention_dim,\n attn_num_head_channels=attn_num_head_channels,\n dual_cross_attention=dual_cross_attention,\n use_linear_projection=use_linear_projection,\n only_cross_attention=only_cross_attention,\n upcast_attention=upcast_attention,\n resnet_time_scale_shift=resnet_time_scale_shift,\n\n unet_use_cross_frame_attention=unet_use_cross_frame_attention,\n unet_use_temporal_attention=unet_use_temporal_attention,\n use_inflated_groupnorm=use_inflated_groupnorm,\n \n use_motion_module=use_motion_module,\n motion_module_type=motion_module_type,\n motion_module_kwargs=motion_module_kwargs,\n )\n raise ValueError(f\"{down_block_type} does not exist.\")\n\n\ndef get_up_block(\n up_block_type,\n num_layers,\n in_channels,\n out_channels,\n prev_output_channel,\n temb_channels,\n add_upsample,\n resnet_eps,\n resnet_act_fn,\n attn_num_head_channels,\n resnet_groups=None,\n cross_attention_dim=None,\n dual_cross_attention=False,\n use_linear_projection=False,\n only_cross_attention=False,\n upcast_attention=False,\n resnet_time_scale_shift=\"default\",\n\n unet_use_cross_frame_attention=False,\n unet_use_temporal_attention=False,\n use_inflated_groupnorm=False,\n \n use_motion_module=None,\n motion_module_type=None,\n motion_module_kwargs=None,\n):\n up_block_type = up_block_type[7:] if up_block_type.startswith(\"UNetRes\") else up_block_type\n if up_block_type == \"UpBlock3D\":\n return UpBlock3D(\n num_layers=num_layers,\n in_channels=in_channels,\n out_channels=out_channels,\n prev_output_channel=prev_output_channel,\n temb_channels=temb_channels,\n add_upsample=add_upsample,\n resnet_eps=resnet_eps,\n resnet_act_fn=resnet_act_fn,\n resnet_groups=resnet_groups,\n resnet_time_scale_shift=resnet_time_scale_shift,\n\n use_inflated_groupnorm=use_inflated_groupnorm,\n\n use_motion_module=use_motion_module,\n motion_module_type=motion_module_type,\n motion_module_kwargs=motion_module_kwargs,\n )\n elif up_block_type == \"CrossAttnUpBlock3D\":\n if cross_attention_dim is None:\n raise ValueError(\"cross_attention_dim must be specified for CrossAttnUpBlock3D\")\n return CrossAttnUpBlock3D(\n num_layers=num_layers,\n in_channels=in_channels,\n out_channels=out_channels,\n prev_output_channel=prev_output_channel,\n temb_channels=temb_channels,\n add_upsample=add_upsample,\n resnet_eps=resnet_eps,\n resnet_act_fn=resnet_act_fn,\n resnet_groups=resnet_groups,\n cross_attention_dim=cross_attention_dim,\n attn_num_head_channels=attn_num_head_channels,\n dual_cross_attention=dual_cross_attention,\n use_linear_projection=use_linear_projection,\n only_cross_attention=only_cross_attention,\n upcast_attention=upcast_attention,\n resnet_time_scale_shift=resnet_time_scale_shift,\n\n unet_use_cross_frame_attention=unet_use_cross_frame_attention,\n unet_use_temporal_attention=unet_use_temporal_attention,\n use_inflated_groupnorm=use_inflated_groupnorm,\n\n use_motion_module=use_motion_module,\n motion_module_type=motion_module_type,\n motion_module_kwargs=motion_module_kwargs,\n )\n raise ValueError(f\"{up_block_type} does not exist.\")\n\n\nclass UNetMidBlock3DCrossAttn(nn.Module):\n def __init__(\n self,\n in_channels: int,\n temb_channels: int,\n dropout: float = 0.0,\n num_layers: int = 1,\n resnet_eps: float = 1e-6,\n resnet_time_scale_shift: str = \"default\",\n resnet_act_fn: str = \"swish\",\n resnet_groups: int = 32,\n resnet_pre_norm: bool = True,\n attn_num_head_channels=1,\n output_scale_factor=1.0,\n cross_attention_dim=1280,\n dual_cross_attention=False,\n use_linear_projection=False,\n upcast_attention=False,\n\n unet_use_cross_frame_attention=False,\n unet_use_temporal_attention=False,\n use_inflated_groupnorm=False,\n\n use_motion_module=None,\n \n motion_module_type=None,\n motion_module_kwargs=None,\n ):\n super().__init__()\n\n self.has_cross_attention = True\n self.attn_num_head_channels = attn_num_head_channels\n resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)\n\n # there is always at least one resnet\n resnets = [\n ResnetBlock3D(\n in_channels=in_channels,\n out_channels=in_channels,\n temb_channels=temb_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n time_embedding_norm=resnet_time_scale_shift,\n non_linearity=resnet_act_fn,\n output_scale_factor=output_scale_factor,\n pre_norm=resnet_pre_norm,\n\n use_inflated_groupnorm=use_inflated_groupnorm,\n )\n ]\n attentions = []\n motion_modules = []\n\n for _ in range(num_layers):\n if dual_cross_attention:\n raise NotImplementedError\n attentions.append(\n Transformer3DModel(\n attn_num_head_channels,\n in_channels // attn_num_head_channels,\n in_channels=in_channels,\n num_layers=1,\n cross_attention_dim=cross_attention_dim,\n norm_num_groups=resnet_groups,\n use_linear_projection=use_linear_projection,\n upcast_attention=upcast_attention,\n\n unet_use_cross_frame_attention=unet_use_cross_frame_attention,\n unet_use_temporal_attention=unet_use_temporal_attention,\n )\n )\n motion_modules.append(\n get_motion_module(\n in_channels=in_channels,\n motion_module_type=motion_module_type, \n motion_module_kwargs=motion_module_kwargs,\n ) if use_motion_module else None\n )\n resnets.append(\n ResnetBlock3D(\n in_channels=in_channels,\n out_channels=in_channels,\n temb_channels=temb_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n time_embedding_norm=resnet_time_scale_shift,\n non_linearity=resnet_act_fn,\n output_scale_factor=output_scale_factor,\n pre_norm=resnet_pre_norm,\n\n use_inflated_groupnorm=use_inflated_groupnorm,\n )\n )\n\n self.attentions = nn.ModuleList(attentions)\n self.resnets = nn.ModuleList(resnets)\n self.motion_modules = nn.ModuleList(motion_modules)\n\n def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):\n hidden_states = self.resnets[0](hidden_states, temb)\n for attn, resnet, motion_module in zip(self.attentions, self.resnets[1:], self.motion_modules):\n hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample\n hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states) if motion_module is not None else hidden_states\n hidden_states = resnet(hidden_states, temb)\n\n return hidden_states\n\n\nclass CrossAttnDownBlock3D(nn.Module):\n def __init__(\n self,\n in_channels: int,\n out_channels: int,\n temb_channels: int,\n dropout: float = 0.0,\n num_layers: int = 1,\n resnet_eps: float = 1e-6,\n resnet_time_scale_shift: str = \"default\",\n resnet_act_fn: str = \"swish\",\n resnet_groups: int = 32,\n resnet_pre_norm: bool = True,\n attn_num_head_channels=1,\n cross_attention_dim=1280,\n output_scale_factor=1.0,\n downsample_padding=1,\n add_downsample=True,\n dual_cross_attention=False,\n use_linear_projection=False,\n only_cross_attention=False,\n upcast_attention=False,\n\n unet_use_cross_frame_attention=False,\n unet_use_temporal_attention=False,\n use_inflated_groupnorm=False,\n \n use_motion_module=None,\n\n motion_module_type=None,\n motion_module_kwargs=None,\n ):\n super().__init__()\n resnets = []\n attentions = []\n motion_modules = []\n\n self.has_cross_attention = True\n self.attn_num_head_channels = attn_num_head_channels\n\n for i in range(num_layers):\n in_channels = in_channels if i == 0 else out_channels\n resnets.append(\n ResnetBlock3D(\n in_channels=in_channels,\n out_channels=out_channels,\n temb_channels=temb_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n time_embedding_norm=resnet_time_scale_shift,\n non_linearity=resnet_act_fn,\n output_scale_factor=output_scale_factor,\n pre_norm=resnet_pre_norm,\n\n use_inflated_groupnorm=use_inflated_groupnorm,\n )\n )\n if dual_cross_attention:\n raise NotImplementedError\n attentions.append(\n Transformer3DModel(\n attn_num_head_channels,\n out_channels // attn_num_head_channels,\n in_channels=out_channels,\n num_layers=1,\n cross_attention_dim=cross_attention_dim,\n norm_num_groups=resnet_groups,\n use_linear_projection=use_linear_projection,\n only_cross_attention=only_cross_attention,\n upcast_attention=upcast_attention,\n\n unet_use_cross_frame_attention=unet_use_cross_frame_attention,\n unet_use_temporal_attention=unet_use_temporal_attention,\n )\n )\n motion_modules.append(\n get_motion_module(\n in_channels=out_channels,\n motion_module_type=motion_module_type, \n motion_module_kwargs=motion_module_kwargs,\n ) if use_motion_module else None\n )\n \n self.attentions = nn.ModuleList(attentions)\n self.resnets = nn.ModuleList(resnets)\n self.motion_modules = nn.ModuleList(motion_modules)\n\n if add_downsample:\n self.downsamplers = nn.ModuleList(\n [\n Downsample3D(\n out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name=\"op\"\n )\n ]\n )\n else:\n self.downsamplers = None\n\n self.gradient_checkpointing = False\n\n def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):\n output_states = ()\n\n for resnet, attn, motion_module in zip(self.resnets, self.attentions, self.motion_modules):\n if self.training and self.gradient_checkpointing:\n\n def create_custom_forward(module, return_dict=None):\n def custom_forward(*inputs):\n if return_dict is not None:\n return module(*inputs, return_dict=return_dict)\n else:\n return module(*inputs)\n\n return custom_forward\n\n hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)\n hidden_states = torch.utils.checkpoint.checkpoint(\n create_custom_forward(attn, return_dict=False),\n hidden_states,\n encoder_hidden_states,\n )[0]\n if motion_module is not None:\n hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(motion_module), hidden_states.requires_grad_(), temb, encoder_hidden_states)\n \n else:\n hidden_states = resnet(hidden_states, temb)\n hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample\n \n # add motion module\n hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states) if motion_module is not None else hidden_states\n\n output_states += (hidden_states,)\n\n if self.downsamplers is not None:\n for downsampler in self.downsamplers:\n hidden_states = downsampler(hidden_states)\n\n output_states += (hidden_states,)\n\n return hidden_states, output_states\n\n\nclass DownBlock3D(nn.Module):\n def __init__(\n self,\n in_channels: int,\n out_channels: int,\n temb_channels: int,\n dropout: float = 0.0,\n num_layers: int = 1,\n resnet_eps: float = 1e-6,\n resnet_time_scale_shift: str = \"default\",\n resnet_act_fn: str = \"swish\",\n resnet_groups: int = 32,\n resnet_pre_norm: bool = True,\n output_scale_factor=1.0,\n add_downsample=True,\n downsample_padding=1,\n\n use_inflated_groupnorm=False,\n \n use_motion_module=None,\n motion_module_type=None,\n motion_module_kwargs=None,\n ):\n super().__init__()\n resnets = []\n motion_modules = []\n\n for i in range(num_layers):\n in_channels = in_channels if i == 0 else out_channels\n resnets.append(\n ResnetBlock3D(\n in_channels=in_channels,\n out_channels=out_channels,\n temb_channels=temb_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n time_embedding_norm=resnet_time_scale_shift,\n non_linearity=resnet_act_fn,\n output_scale_factor=output_scale_factor,\n pre_norm=resnet_pre_norm,\n\n use_inflated_groupnorm=use_inflated_groupnorm,\n )\n )\n motion_modules.append(\n get_motion_module(\n in_channels=out_channels,\n motion_module_type=motion_module_type, \n motion_module_kwargs=motion_module_kwargs,\n ) if use_motion_module else None\n )\n \n self.resnets = nn.ModuleList(resnets)\n self.motion_modules = nn.ModuleList(motion_modules)\n\n if add_downsample:\n self.downsamplers = nn.ModuleList(\n [\n Downsample3D(\n out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name=\"op\"\n )\n ]\n )\n else:\n self.downsamplers = None\n\n self.gradient_checkpointing = False\n\n def forward(self, hidden_states, temb=None, encoder_hidden_states=None):\n output_states = ()\n\n for resnet, motion_module in zip(self.resnets, self.motion_modules):\n if self.training and self.gradient_checkpointing:\n def create_custom_forward(module):\n def custom_forward(*inputs):\n return module(*inputs)\n\n return custom_forward\n\n hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)\n if motion_module is not None:\n hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(motion_module), hidden_states.requires_grad_(), temb, encoder_hidden_states)\n else:\n hidden_states = resnet(hidden_states, temb)\n\n # add motion module\n hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states) if motion_module is not None else hidden_states\n\n output_states += (hidden_states,)\n\n if self.downsamplers is not None:\n for downsampler in self.downsamplers:\n hidden_states = downsampler(hidden_states)\n\n output_states += (hidden_states,)\n\n return hidden_states, output_states\n\n\nclass CrossAttnUpBlock3D(nn.Module):\n def __init__(\n self,\n in_channels: int,\n out_channels: int,\n prev_output_channel: int,\n temb_channels: int,\n dropout: float = 0.0,\n num_layers: int = 1,\n resnet_eps: float = 1e-6,\n resnet_time_scale_shift: str = \"default\",\n resnet_act_fn: str = \"swish\",\n resnet_groups: int = 32,\n resnet_pre_norm: bool = True,\n attn_num_head_channels=1,\n cross_attention_dim=1280,\n output_scale_factor=1.0,\n add_upsample=True,\n dual_cross_attention=False,\n use_linear_projection=False,\n only_cross_attention=False,\n upcast_attention=False,\n\n unet_use_cross_frame_attention=False,\n unet_use_temporal_attention=False,\n use_inflated_groupnorm=False,\n \n use_motion_module=None,\n\n motion_module_type=None,\n motion_module_kwargs=None,\n ):\n super().__init__()\n resnets = []\n attentions = []\n motion_modules = []\n\n self.has_cross_attention = True\n self.attn_num_head_channels = attn_num_head_channels\n\n for i in range(num_layers):\n res_skip_channels = in_channels if (i == num_layers - 1) else out_channels\n resnet_in_channels = prev_output_channel if i == 0 else out_channels\n\n resnets.append(\n ResnetBlock3D(\n in_channels=resnet_in_channels + res_skip_channels,\n out_channels=out_channels,\n temb_channels=temb_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n time_embedding_norm=resnet_time_scale_shift,\n non_linearity=resnet_act_fn,\n output_scale_factor=output_scale_factor,\n pre_norm=resnet_pre_norm,\n\n use_inflated_groupnorm=use_inflated_groupnorm,\n )\n )\n if dual_cross_attention:\n raise NotImplementedError\n attentions.append(\n Transformer3DModel(\n attn_num_head_channels,\n out_channels // attn_num_head_channels,\n in_channels=out_channels,\n num_layers=1,\n cross_attention_dim=cross_attention_dim,\n norm_num_groups=resnet_groups,\n use_linear_projection=use_linear_projection,\n only_cross_attention=only_cross_attention,\n upcast_attention=upcast_attention,\n\n unet_use_cross_frame_attention=unet_use_cross_frame_attention,\n unet_use_temporal_attention=unet_use_temporal_attention,\n )\n )\n motion_modules.append(\n get_motion_module(\n in_channels=out_channels,\n motion_module_type=motion_module_type, \n motion_module_kwargs=motion_module_kwargs,\n ) if use_motion_module else None\n )\n \n self.attentions = nn.ModuleList(attentions)\n self.resnets = nn.ModuleList(resnets)\n self.motion_modules = nn.ModuleList(motion_modules)\n\n if add_upsample:\n self.upsamplers = nn.ModuleList([Upsample3D(out_channels, use_conv=True, out_channels=out_channels)])\n else:\n self.upsamplers = None\n\n self.gradient_checkpointing = False\n\n def forward(\n self,\n hidden_states,\n res_hidden_states_tuple,\n temb=None,\n encoder_hidden_states=None,\n upsample_size=None,\n attention_mask=None,\n ):\n for resnet, attn, motion_module in zip(self.resnets, self.attentions, self.motion_modules):\n # pop res hidden states\n res_hidden_states = res_hidden_states_tuple[-1]\n res_hidden_states_tuple = res_hidden_states_tuple[:-1]\n hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)\n\n if self.training and self.gradient_checkpointing:\n\n def create_custom_forward(module, return_dict=None):\n def custom_forward(*inputs):\n if return_dict is not None:\n return module(*inputs, return_dict=return_dict)\n else:\n return module(*inputs)\n\n return custom_forward\n\n hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)\n hidden_states = torch.utils.checkpoint.checkpoint(\n create_custom_forward(attn, return_dict=False),\n hidden_states,\n encoder_hidden_states,\n )[0]\n if motion_module is not None:\n hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(motion_module), hidden_states.requires_grad_(), temb, encoder_hidden_states)\n \n else:\n hidden_states = resnet(hidden_states, temb)\n hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample\n \n # add motion module\n hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states) if motion_module is not None else hidden_states\n\n if self.upsamplers is not None:\n for upsampler in self.upsamplers:\n hidden_states = upsampler(hidden_states, upsample_size)\n\n return hidden_states\n\n\nclass UpBlock3D(nn.Module):\n def __init__(\n self,\n in_channels: int,\n prev_output_channel: int,\n out_channels: int,\n temb_channels: int,\n dropout: float = 0.0,\n num_layers: int = 1,\n resnet_eps: float = 1e-6,\n resnet_time_scale_shift: str = \"default\",\n resnet_act_fn: str = \"swish\",\n resnet_groups: int = 32,\n resnet_pre_norm: bool = True,\n output_scale_factor=1.0,\n add_upsample=True,\n\n use_inflated_groupnorm=False,\n\n use_motion_module=None,\n motion_module_type=None,\n motion_module_kwargs=None,\n ):\n super().__init__()\n resnets = []\n motion_modules = []\n\n for i in range(num_layers):\n res_skip_channels = in_channels if (i == num_layers - 1) else out_channels\n resnet_in_channels = prev_output_channel if i == 0 else out_channels\n\n resnets.append(\n ResnetBlock3D(\n in_channels=resnet_in_channels + res_skip_channels,\n out_channels=out_channels,\n temb_channels=temb_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n time_embedding_norm=resnet_time_scale_shift,\n non_linearity=resnet_act_fn,\n output_scale_factor=output_scale_factor,\n pre_norm=resnet_pre_norm,\n\n use_inflated_groupnorm=use_inflated_groupnorm,\n )\n )\n motion_modules.append(\n get_motion_module(\n in_channels=out_channels,\n motion_module_type=motion_module_type, \n motion_module_kwargs=motion_module_kwargs,\n ) if use_motion_module else None\n )\n\n self.resnets = nn.ModuleList(resnets)\n self.motion_modules = nn.ModuleList(motion_modules)\n\n if add_upsample:\n self.upsamplers = nn.ModuleList([Upsample3D(out_channels, use_conv=True, out_channels=out_channels)])\n else:\n self.upsamplers = None\n\n self.gradient_checkpointing = False\n\n def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, encoder_hidden_states=None,):\n for resnet, motion_module in zip(self.resnets, self.motion_modules):\n # pop res hidden states\n res_hidden_states = res_hidden_states_tuple[-1]\n res_hidden_states_tuple = res_hidden_states_tuple[:-1]\n hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)\n\n if self.training and self.gradient_checkpointing:\n def create_custom_forward(module):\n def custom_forward(*inputs):\n return module(*inputs)\n\n return custom_forward\n\n hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)\n if motion_module is not None:\n hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(motion_module), hidden_states.requires_grad_(), temb, encoder_hidden_states)\n else:\n hidden_states = resnet(hidden_states, temb)\n hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states) if motion_module is not None else hidden_states\n\n if self.upsamplers is not None:\n for upsampler in self.upsamplers:\n hidden_states = upsampler(hidden_states, upsample_size)\n\n return hidden_states\n" }, { "path": "motionclone/pipelines/pipeline_animation.py", "content": "# Adapted from https://github.com/showlab/Tune-A-Video/blob/main/tuneavideo/pipelines/pipeline_tuneavideo.py\n\nimport inspect\nfrom typing import Callable, List, Optional, Union, Any, Dict\nfrom dataclasses import dataclass\nfrom diffusers import StableDiffusionPipeline, DDIMInverseScheduler\n\nimport os\nimport pickle\nimport numpy as np\nimport torch\nfrom tqdm import tqdm\nimport omegaconf\nfrom omegaconf import OmegaConf\nimport yaml\nfrom diffusers.utils import is_accelerate_available\nfrom packaging import version\nfrom transformers import CLIPTextModel, CLIPTokenizer\nfrom diffusers.configuration_utils import FrozenDict\nfrom diffusers.models import AutoencoderKL\nfrom diffusers.pipeline_utils import DiffusionPipeline\nfrom diffusers.schedulers import (\n DDIMScheduler,\n DPMSolverMultistepScheduler,\n EulerAncestralDiscreteScheduler,\n EulerDiscreteScheduler,\n LMSDiscreteScheduler,\n PNDMScheduler,\n)\nfrom diffusers.utils import deprecate, logging, BaseOutput\nfrom einops import rearrange\nfrom ..models.unet import UNet3DConditionModel\nfrom ..models.sparse_controlnet import SparseControlNetModel\nfrom ..utils.xformer_attention import *\nfrom ..utils.conv_layer import *\nfrom ..utils.util import *\n\nlogger = logging.get_logger(__name__) # pylint: disable=invalid-name\n\n\n@dataclass\nclass AnimationPipelineOutput(BaseOutput):\n videos: Union[torch.Tensor, np.ndarray]\n\n\nclass AnimationPipeline(DiffusionPipeline):\n _optional_components = []\n\n def __init__(\n self,\n vae: AutoencoderKL,\n text_encoder: CLIPTextModel,\n tokenizer: CLIPTokenizer,\n unet: UNet3DConditionModel,\n scheduler: Union[\n DDIMScheduler,\n PNDMScheduler,\n LMSDiscreteScheduler,\n EulerDiscreteScheduler,\n EulerAncestralDiscreteScheduler,\n DPMSolverMultistepScheduler,\n ],\n controlnet: Union[SparseControlNetModel, None] = None,\n ):\n super().__init__()\n\n if hasattr(scheduler.config, \"steps_offset\") and scheduler.config.steps_offset != 1:\n deprecation_message = (\n f\"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`\"\n f\" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure \"\n \"to update the config accordingly as leaving `steps_offset` might led to incorrect results\"\n \" in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,\"\n \" it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`\"\n \" file\"\n )\n deprecate(\"steps_offset!=1\", \"1.0.0\", deprecation_message, standard_warn=False)\n new_config = dict(scheduler.config)\n new_config[\"steps_offset\"] = 1\n scheduler._internal_dict = FrozenDict(new_config)\n\n if hasattr(scheduler.config, \"clip_sample\") and scheduler.config.clip_sample is True:\n deprecation_message = (\n f\"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`.\"\n \" `clip_sample` should be set to False in the configuration file. Please make sure to update the\"\n \" config accordingly as not setting `clip_sample` in the config might lead to incorrect results in\"\n \" future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very\"\n \" nice if you could open a Pull request for the `scheduler/scheduler_config.json` file\"\n )\n deprecate(\"clip_sample not set\", \"1.0.0\", deprecation_message, standard_warn=False)\n new_config = dict(scheduler.config)\n new_config[\"clip_sample\"] = False\n scheduler._internal_dict = FrozenDict(new_config)\n\n is_unet_version_less_0_9_0 = hasattr(unet.config, \"_diffusers_version\") and version.parse(\n version.parse(unet.config._diffusers_version).base_version\n ) < version.parse(\"0.9.0.dev0\")\n is_unet_sample_size_less_64 = hasattr(unet.config, \"sample_size\") and unet.config.sample_size < 64\n if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:\n deprecation_message = (\n \"The configuration file of the unet has set the default `sample_size` to smaller than\"\n \" 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the\"\n \" following: \\n- CompVis/stable-diffusion-v1-4 \\n- CompVis/stable-diffusion-v1-3 \\n-\"\n \" CompVis/stable-diffusion-v1-2 \\n- CompVis/stable-diffusion-v1-1 \\n- runwayml/stable-diffusion-v1-5\"\n \" \\n- runwayml/stable-diffusion-inpainting \\n you should change 'sample_size' to 64 in the\"\n \" configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`\"\n \" in the config might lead to incorrect results in future versions. If you have downloaded this\"\n \" checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for\"\n \" the `unet/config.json` file\"\n )\n deprecate(\"sample_size<64\", \"1.0.0\", deprecation_message, standard_warn=False)\n new_config = dict(unet.config)\n new_config[\"sample_size\"] = 64\n unet._internal_dict = FrozenDict(new_config)\n\n self.register_modules(\n vae=vae,\n text_encoder=text_encoder,\n tokenizer=tokenizer,\n unet=unet,\n scheduler=scheduler,\n controlnet=controlnet,\n )\n self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)\n\n \n\n def enable_vae_slicing(self):\n self.vae.enable_slicing()\n\n def disable_vae_slicing(self):\n self.vae.disable_slicing()\n\n def enable_sequential_cpu_offload(self, gpu_id=0):\n if is_accelerate_available():\n from accelerate import cpu_offload\n else:\n raise ImportError(\"Please install accelerate via `pip install accelerate`\")\n\n device = torch.device(f\"cuda:{gpu_id}\")\n\n for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:\n if cpu_offloaded_model is not None:\n cpu_offload(cpu_offloaded_model, device)\n \n\n \n @property\n def _execution_device(self):\n if self.device != torch.device(\"meta\") or not hasattr(self.unet, \"_hf_hook\"):\n return self.device\n for module in self.unet.modules():\n if (\n hasattr(module, \"_hf_hook\")\n and hasattr(module._hf_hook, \"execution_device\")\n and module._hf_hook.execution_device is not None\n ):\n return torch.device(module._hf_hook.execution_device)\n return self.device\n\n def _encode_prompt(self, prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt):\n batch_size = len(prompt) if isinstance(prompt, list) else 1\n \n text_inputs = self.tokenizer(\n prompt,\n padding=\"max_length\",\n max_length=self.tokenizer.model_max_length,\n truncation=True,\n return_tensors=\"pt\",\n )\n text_input_ids = text_inputs.input_ids\n untruncated_ids = self.tokenizer(prompt, padding=\"longest\", return_tensors=\"pt\").input_ids\n\n if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):\n removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])\n logger.warning(\n \"The following part of your input was truncated because CLIP can only handle sequences up to\"\n f\" {self.tokenizer.model_max_length} tokens: {removed_text}\"\n )\n\n if hasattr(self.text_encoder.config, \"use_attention_mask\") and self.text_encoder.config.use_attention_mask:\n attention_mask = text_inputs.attention_mask.to(device)\n else:\n attention_mask = None\n\n text_embeddings = self.text_encoder(\n text_input_ids.to(device),\n attention_mask=attention_mask,\n )\n text_embeddings = text_embeddings[0]\n\n # duplicate text embeddings for each generation per prompt, using mps friendly method\n bs_embed, seq_len, _ = text_embeddings.shape\n text_embeddings = text_embeddings.repeat(1, num_videos_per_prompt, 1)\n text_embeddings = text_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1)\n\n # get unconditional embeddings for classifier free guidance\n if do_classifier_free_guidance:\n uncond_tokens: List[str]\n if negative_prompt is None:\n uncond_tokens = [\"\"] * batch_size\n elif type(prompt) is not type(negative_prompt):\n raise TypeError(\n f\"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=\"\n f\" {type(prompt)}.\"\n )\n elif isinstance(negative_prompt, str):\n uncond_tokens = [negative_prompt]\n elif batch_size != len(negative_prompt):\n raise ValueError(\n f\"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:\"\n f\" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches\"\n \" the batch size of `prompt`.\"\n )\n else:\n uncond_tokens = negative_prompt\n\n max_length = text_input_ids.shape[-1]\n uncond_input = self.tokenizer(\n uncond_tokens,\n padding=\"max_length\",\n max_length=max_length,\n truncation=True,\n return_tensors=\"pt\",\n )\n\n if hasattr(self.text_encoder.config, \"use_attention_mask\") and self.text_encoder.config.use_attention_mask:\n attention_mask = uncond_input.attention_mask.to(device)\n else:\n attention_mask = None\n\n uncond_embeddings = self.text_encoder(\n uncond_input.input_ids.to(device),\n attention_mask=attention_mask,\n )\n uncond_embeddings = uncond_embeddings[0]\n\n # duplicate unconditional embeddings for each generation per prompt, using mps friendly method\n seq_len = uncond_embeddings.shape[1]\n uncond_embeddings = uncond_embeddings.repeat(1, num_videos_per_prompt, 1)\n uncond_embeddings = uncond_embeddings.view(batch_size * num_videos_per_prompt, seq_len, -1)\n\n # For classifier free guidance, we need to do two forward passes.\n # Here we concatenate the unconditional and text embeddings into a single batch\n # to avoid doing two forward passes\n text_embeddings = torch.cat([uncond_embeddings, text_embeddings])\n\n return text_embeddings\n\n @torch.no_grad()\n def decode_latents(self, latents):\n video_length = latents.shape[2]\n latents = 1 / 0.18215 * latents\n latents = rearrange(latents, \"b c f h w -> (b f) c h w\")\n # video = self.vae.decode(latents).sample\n video = []\n for frame_idx in tqdm(range(latents.shape[0])):\n video.append(self.vae.decode(latents[frame_idx:frame_idx+1]).sample)\n video = torch.cat(video)\n video = rearrange(video, \"(b f) c h w -> b c f h w\", f=video_length)\n video = (video / 2 + 0.5).clamp(0, 1)\n # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16\n video = video.cpu().float().numpy()\n return video\n\n def prepare_extra_step_kwargs(self, generator, eta):\n # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature\n # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.\n # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502\n # and should be between [0, 1]\n\n accepts_eta = \"eta\" in set(inspect.signature(self.scheduler.step).parameters.keys())\n extra_step_kwargs = {}\n if accepts_eta:\n extra_step_kwargs[\"eta\"] = eta\n\n # check if the scheduler accepts generator\n accepts_generator = \"generator\" in set(inspect.signature(self.scheduler.step).parameters.keys())\n if accepts_generator:\n extra_step_kwargs[\"generator\"] = generator\n return extra_step_kwargs\n\n def check_inputs(self, prompt, height, width, callback_steps):\n if not isinstance(prompt, str) and not isinstance(prompt, list):\n raise ValueError(f\"`prompt` has to be of type `str` or `list` but is {type(prompt)}\")\n\n if height % 8 != 0 or width % 8 != 0:\n raise ValueError(f\"`height` and `width` have to be divisible by 8 but are {height} and {width}.\")\n\n if (callback_steps is None) or (\n callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)\n ):\n raise ValueError(\n f\"`callback_steps` has to be a positive integer but is {callback_steps} of type\"\n f\" {type(callback_steps)}.\"\n )\n\n def prepare_latents(self, batch_size, num_channels_latents, video_length, height, width, dtype, device, generator, latents=None):\n shape = (batch_size, num_channels_latents, video_length, height // self.vae_scale_factor, width // self.vae_scale_factor)\n if isinstance(generator, list) and len(generator) != batch_size:\n raise ValueError(\n f\"You have passed a list of generators of length {len(generator)}, but requested an effective batch\"\n f\" size of {batch_size}. Make sure the batch size matches the length of the generators.\"\n )\n if latents is None:\n rand_device = \"cpu\" if device.type == \"mps\" else device\n\n if isinstance(generator, list):\n shape = shape\n # shape = (1,) + shape[1:]\n latents = [\n torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype)\n for i in range(batch_size)\n ]\n latents = torch.cat(latents, dim=0).to(device)\n else:\n latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device)\n else:\n if latents.shape != shape:\n raise ValueError(f\"Unexpected latents shape, got {latents.shape}, expected {shape}\")\n latents = latents.to(device)\n\n # scale the initial noise by the standard deviation required by the scheduler\n latents = latents * self.scheduler.init_noise_sigma\n return latents\n\n @torch.no_grad()\n def __call__(\n self,\n prompt: Union[str, List[str]],\n video_length: Optional[int],\n height: Optional[int] = None,\n width: Optional[int] = None,\n num_inference_steps: int = 50,\n guidance_scale: float = 7.5,\n negative_prompt: Optional[Union[str, List[str]]] = None,\n num_videos_per_prompt: Optional[int] = 1,\n eta: float = 0.0,\n generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n latents: Optional[torch.FloatTensor] = None,\n output_type: Optional[str] = \"tensor\",\n return_dict: bool = True,\n callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,\n callback_steps: Optional[int] = 1,\n save_probs = False,\n # support controlnet\n controlnet_images: torch.FloatTensor = None,\n controlnet_image_index: list = [0],\n controlnet_conditioning_scale: Union[float, List[float]] = 1.0,\n\n **kwargs,\n ):\n\n # to record temp attention probs\n self.unet = prep_unet_attention(self.unet)\n self.unet = prep_unet_conv(self.unet)\n self.guidance_config = guidance_scale\n\n self.temp_rec = {}\n\n # Default height and width to unet\n height = height or self.unet.config.sample_size * self.vae_scale_factor\n width = width or self.unet.config.sample_size * self.vae_scale_factor\n\n # Check inputs. Raise error if not correct\n self.check_inputs(prompt, height, width, callback_steps)\n\n # Define call parameters\n # batch_size = 1 if isinstance(prompt, str) else len(prompt)\n batch_size = 1\n if latents is not None:\n batch_size = latents.shape[0]\n if isinstance(prompt, list):\n batch_size = len(prompt)\n\n device = self._execution_device\n # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)\n # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`\n # corresponds to doing no classifier free guidance.\n do_classifier_free_guidance = guidance_scale > 1.0\n \n # Encode input prompt\n prompt = prompt if isinstance(prompt, list) else [prompt] * batch_size\n if negative_prompt is not None:\n negative_prompt = negative_prompt if isinstance(negative_prompt, list) else [negative_prompt] * batch_size \n text_embeddings = self._encode_prompt(\n prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt\n )\n\n # import pdb; pdb.set_trace()\n # Prepare timesteps\n self.scheduler.set_timesteps(num_inference_steps, device=device)\n timesteps = self.scheduler.timesteps\n\n # Prepare latent variables\n num_channels_latents = self.unet.in_channels\n latents = self.prepare_latents(\n batch_size * num_videos_per_prompt,\n num_channels_latents,\n video_length,\n height,\n width,\n text_embeddings.dtype,\n device,\n generator,\n latents,\n )\n latents_dtype = latents.dtype\n\n # Prepare extra step kwargs.\n extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)\n\n # Denoising loop\n num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order\n with self.progress_bar(total=num_inference_steps) as progress_bar:\n for i, t in enumerate(timesteps):\n # expand the latents if we are doing classifier free guidance\n latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n\n down_block_additional_residuals = mid_block_additional_residual = None\n # import pdb; pdb.set_trace()\n if (getattr(self, \"controlnet\", None) != None) and (controlnet_images != None):\n assert controlnet_images.dim() == 5\n\n controlnet_noisy_latents = latent_model_input\n controlnet_prompt_embeds = text_embeddings\n\n controlnet_images = controlnet_images.to(latents.device)\n\n controlnet_cond_shape = list(controlnet_images.shape)\n controlnet_cond_shape[2] = video_length\n controlnet_cond = torch.zeros(controlnet_cond_shape).to(latents.device)\n\n controlnet_conditioning_mask_shape = list(controlnet_cond.shape)\n controlnet_conditioning_mask_shape[1] = 1\n controlnet_conditioning_mask = torch.zeros(controlnet_conditioning_mask_shape).to(latents.device)\n\n assert controlnet_images.shape[2] >= len(controlnet_image_index)\n controlnet_cond[:,:,controlnet_image_index] = controlnet_images[:,:,:len(controlnet_image_index)]\n controlnet_conditioning_mask[:,:,controlnet_image_index] = 1\n\n down_block_additional_residuals, mid_block_additional_residual = self.controlnet(\n controlnet_noisy_latents, t,\n encoder_hidden_states=controlnet_prompt_embeds,\n controlnet_cond=controlnet_cond,\n conditioning_mask=controlnet_conditioning_mask,\n conditioning_scale=controlnet_conditioning_scale,\n guess_mode=False, return_dict=False,\n )\n \n # predict the noise residual\n noise_pred = self.unet(\n latent_model_input, t, \n encoder_hidden_states=text_embeddings,\n down_block_additional_residuals = down_block_additional_residuals,\n mid_block_additional_residual = mid_block_additional_residual,\n ).sample.to(dtype=latents_dtype)\n\n\n # get temp attn probs\n if save_probs:\n temp_attn_prob = self.get_temp_attn_prob()\n for key in temp_attn_prob.keys():\n temp_attn_prob[key] = temp_attn_prob[key].chunk(2, dim = 0)[0].detach().clone().cpu()\n self.temp_rec[i] = temp_attn_prob\n # import pdb; pdb.set_trace()\n\n # perform guidance\n if do_classifier_free_guidance:\n noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n\n # compute the previous noisy sample x_t -> x_t-1\n latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample\n\n # call the callback, if provided\n if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):\n progress_bar.update()\n if callback is not None and i % callback_steps == 0:\n callback(i, t, latents)\n\n # pickle temp attn prob\n if save_probs:\n with open('temp_dic.pkl', 'wb') as f:\n pickle.dump(self.temp_rec, f)\n\n # Post-processing\n video = self.decode_latents(latents)\n\n # Convert to tensor\n if output_type == \"tensor\":\n video = torch.from_numpy(video)\n\n if not return_dict:\n return video\n\n return AnimationPipelineOutput(videos=video)\n" }, { "path": "motionclone/utils/conv_layer.py", "content": "import torch\n\ndef conv_forward(self):\n def forward(input_tensor, temb, scale=1.0):\n hidden_states = input_tensor\n\n hidden_states = self.norm1(hidden_states)\n hidden_states = self.nonlinearity(hidden_states)\n # import pdb; pdb.set_trace()\n if self.upsample is not None:\n # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984\n if hidden_states.shape[0] >= 64:\n input_tensor = input_tensor.contiguous()\n hidden_states = hidden_states.contiguous()\n input_tensor = self.upsample(input_tensor)\n hidden_states = self.upsample(hidden_states)\n elif self.downsample is not None:\n input_tensor = self.downsample(input_tensor)\n hidden_states = self.downsample(hidden_states)\n\n hidden_states = self.conv1(hidden_states)\n \n if temb is not None:\n temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None, None].repeat(1, 1, hidden_states.shape[2], 1, 1)\n \n if temb is not None and self.time_embedding_norm == \"default\":\n hidden_states = hidden_states + temb\n\n hidden_states = self.norm2(hidden_states)\n\n if temb is not None and self.time_embedding_norm == \"scale_shift\":\n scale, shift = torch.chunk(temb, 2, dim=1)\n hidden_states = hidden_states * (1 + scale) + shift\n\n hidden_states = self.nonlinearity(hidden_states)\n\n hidden_states = self.dropout(hidden_states)\n hidden_states = self.conv2(hidden_states)\n\n # record hidden state\n self.record_hidden_state = hidden_states\n\n if self.conv_shortcut is not None:\n input_tensor = self.conv_shortcut(input_tensor)\n\n output_tensor = (input_tensor + hidden_states) / self.output_scale_factor\n\n return output_tensor\n\n return forward\n\n\ndef get_conv_feat(unet):\n hidden_state_dict = dict()\n for i in range(len(unet.up_blocks)):\n for j in range(len(unet.up_blocks[i].resnets)):\n module = unet.up_blocks[i].resnets[j]\n module_name = f\"up_blocks.{i}.resnets.{j}\"\n # print(module_name)\n hidden_state_dict[module_name] = module.record_hidden_state\n return hidden_state_dict\n\n\ndef prep_unet_conv(unet):\n for i in range(len(unet.up_blocks)):\n for j in range(len(unet.up_blocks[i].resnets)):\n module = unet.up_blocks[i].resnets[j]\n module.forward = conv_forward(module)\n return unet\n" }, { "path": "motionclone/utils/convert_from_ckpt.py", "content": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\" Conversion script for the Stable Diffusion checkpoints.\"\"\"\n\nimport re\nfrom io import BytesIO\nfrom typing import Optional\n\nimport requests\nimport torch\nfrom transformers import (\n AutoFeatureExtractor,\n BertTokenizerFast,\n CLIPImageProcessor,\n CLIPTextModel,\n CLIPTextModelWithProjection,\n CLIPTokenizer,\n CLIPVisionConfig,\n CLIPVisionModelWithProjection,\n)\n\nfrom diffusers.models import (\n AutoencoderKL,\n PriorTransformer,\n UNet2DConditionModel,\n)\nfrom diffusers.schedulers import (\n DDIMScheduler,\n DDPMScheduler,\n DPMSolverMultistepScheduler,\n EulerAncestralDiscreteScheduler,\n EulerDiscreteScheduler,\n HeunDiscreteScheduler,\n LMSDiscreteScheduler,\n PNDMScheduler,\n UnCLIPScheduler,\n)\nfrom diffusers.utils.import_utils import BACKENDS_MAPPING\n\n\ndef shave_segments(path, n_shave_prefix_segments=1):\n \"\"\"\n Removes segments. Positive values shave the first segments, negative shave the last segments.\n \"\"\"\n if n_shave_prefix_segments >= 0:\n return \".\".join(path.split(\".\")[n_shave_prefix_segments:])\n else:\n return \".\".join(path.split(\".\")[:n_shave_prefix_segments])\n\n\ndef renew_resnet_paths(old_list, n_shave_prefix_segments=0):\n \"\"\"\n Updates paths inside resnets to the new naming scheme (local renaming)\n \"\"\"\n mapping = []\n for old_item in old_list:\n new_item = old_item.replace(\"in_layers.0\", \"norm1\")\n new_item = new_item.replace(\"in_layers.2\", \"conv1\")\n\n new_item = new_item.replace(\"out_layers.0\", \"norm2\")\n new_item = new_item.replace(\"out_layers.3\", \"conv2\")\n\n new_item = new_item.replace(\"emb_layers.1\", \"time_emb_proj\")\n new_item = new_item.replace(\"skip_connection\", \"conv_shortcut\")\n\n new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)\n\n mapping.append({\"old\": old_item, \"new\": new_item})\n\n return mapping\n\n\ndef renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):\n \"\"\"\n Updates paths inside resnets to the new naming scheme (local renaming)\n \"\"\"\n mapping = []\n for old_item in old_list:\n new_item = old_item\n\n new_item = new_item.replace(\"nin_shortcut\", \"conv_shortcut\")\n new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)\n\n mapping.append({\"old\": old_item, \"new\": new_item})\n\n return mapping\n\n\ndef renew_attention_paths(old_list, n_shave_prefix_segments=0):\n \"\"\"\n Updates paths inside attentions to the new naming scheme (local renaming)\n \"\"\"\n mapping = []\n for old_item in old_list:\n new_item = old_item\n\n # new_item = new_item.replace('norm.weight', 'group_norm.weight')\n # new_item = new_item.replace('norm.bias', 'group_norm.bias')\n\n # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')\n # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')\n\n # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)\n\n mapping.append({\"old\": old_item, \"new\": new_item})\n\n return mapping\n\n\ndef renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):\n \"\"\"\n Updates paths inside attentions to the new naming scheme (local renaming)\n \"\"\"\n mapping = []\n for old_item in old_list:\n new_item = old_item\n\n new_item = new_item.replace(\"norm.weight\", \"group_norm.weight\")\n new_item = new_item.replace(\"norm.bias\", \"group_norm.bias\")\n\n new_item = new_item.replace(\"q.weight\", \"query.weight\")\n new_item = new_item.replace(\"q.bias\", \"query.bias\")\n\n new_item = new_item.replace(\"k.weight\", \"key.weight\")\n new_item = new_item.replace(\"k.bias\", \"key.bias\")\n\n new_item = new_item.replace(\"v.weight\", \"value.weight\")\n new_item = new_item.replace(\"v.bias\", \"value.bias\")\n\n new_item = new_item.replace(\"proj_out.weight\", \"proj_attn.weight\")\n new_item = new_item.replace(\"proj_out.bias\", \"proj_attn.bias\")\n\n new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)\n\n mapping.append({\"old\": old_item, \"new\": new_item})\n\n return mapping\n\n\ndef assign_to_checkpoint(\n paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None\n):\n \"\"\"\n This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits\n attention layers, and takes into account additional replacements that may arise.\n\n Assigns the weights to the new checkpoint.\n \"\"\"\n assert isinstance(paths, list), \"Paths should be a list of dicts containing 'old' and 'new' keys.\"\n\n # Splits the attention layers into three variables.\n if attention_paths_to_split is not None:\n for path, path_map in attention_paths_to_split.items():\n old_tensor = old_checkpoint[path]\n channels = old_tensor.shape[0] // 3\n\n target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)\n\n num_heads = old_tensor.shape[0] // config[\"num_head_channels\"] // 3\n\n old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])\n query, key, value = old_tensor.split(channels // num_heads, dim=1)\n\n checkpoint[path_map[\"query\"]] = query.reshape(target_shape)\n checkpoint[path_map[\"key\"]] = key.reshape(target_shape)\n checkpoint[path_map[\"value\"]] = value.reshape(target_shape)\n\n for path in paths:\n new_path = path[\"new\"]\n\n # These have already been assigned\n if attention_paths_to_split is not None and new_path in attention_paths_to_split:\n continue\n\n # Global renaming happens here\n new_path = new_path.replace(\"middle_block.0\", \"mid_block.resnets.0\")\n new_path = new_path.replace(\"middle_block.1\", \"mid_block.attentions.0\")\n new_path = new_path.replace(\"middle_block.2\", \"mid_block.resnets.1\")\n\n if additional_replacements is not None:\n for replacement in additional_replacements:\n new_path = new_path.replace(replacement[\"old\"], replacement[\"new\"])\n\n # proj_attn.weight has to be converted from conv 1D to linear\n if \"proj_attn.weight\" in new_path:\n checkpoint[new_path] = old_checkpoint[path[\"old\"]][:, :, 0]\n else:\n checkpoint[new_path] = old_checkpoint[path[\"old\"]]\n\n\ndef conv_attn_to_linear(checkpoint):\n keys = list(checkpoint.keys())\n attn_keys = [\"query.weight\", \"key.weight\", \"value.weight\"]\n for key in keys:\n if \".\".join(key.split(\".\")[-2:]) in attn_keys:\n if checkpoint[key].ndim > 2:\n checkpoint[key] = checkpoint[key][:, :, 0, 0]\n elif \"proj_attn.weight\" in key:\n if checkpoint[key].ndim > 2:\n checkpoint[key] = checkpoint[key][:, :, 0]\n\n\ndef create_unet_diffusers_config(original_config, image_size: int, controlnet=False):\n \"\"\"\n Creates a config for the diffusers based on the config of the LDM model.\n \"\"\"\n if controlnet:\n unet_params = original_config.model.params.control_stage_config.params\n else:\n unet_params = original_config.model.params.unet_config.params\n\n vae_params = original_config.model.params.first_stage_config.params.ddconfig\n\n block_out_channels = [unet_params.model_channels * mult for mult in unet_params.channel_mult]\n\n down_block_types = []\n resolution = 1\n for i in range(len(block_out_channels)):\n block_type = \"CrossAttnDownBlock2D\" if resolution in unet_params.attention_resolutions else \"DownBlock2D\"\n down_block_types.append(block_type)\n if i != len(block_out_channels) - 1:\n resolution *= 2\n\n up_block_types = []\n for i in range(len(block_out_channels)):\n block_type = \"CrossAttnUpBlock2D\" if resolution in unet_params.attention_resolutions else \"UpBlock2D\"\n up_block_types.append(block_type)\n resolution //= 2\n\n vae_scale_factor = 2 ** (len(vae_params.ch_mult) - 1)\n\n head_dim = unet_params.num_heads if \"num_heads\" in unet_params else None\n use_linear_projection = (\n unet_params.use_linear_in_transformer if \"use_linear_in_transformer\" in unet_params else False\n )\n if use_linear_projection:\n # stable diffusion 2-base-512 and 2-768\n if head_dim is None:\n head_dim = [5, 10, 20, 20]\n\n class_embed_type = None\n projection_class_embeddings_input_dim = None\n\n if \"num_classes\" in unet_params:\n if unet_params.num_classes == \"sequential\":\n class_embed_type = \"projection\"\n assert \"adm_in_channels\" in unet_params\n projection_class_embeddings_input_dim = unet_params.adm_in_channels\n else:\n raise NotImplementedError(f\"Unknown conditional unet num_classes config: {unet_params.num_classes}\")\n\n config = {\n \"sample_size\": image_size // vae_scale_factor,\n \"in_channels\": unet_params.in_channels,\n \"down_block_types\": tuple(down_block_types),\n \"block_out_channels\": tuple(block_out_channels),\n \"layers_per_block\": unet_params.num_res_blocks,\n \"cross_attention_dim\": unet_params.context_dim,\n \"attention_head_dim\": head_dim,\n \"use_linear_projection\": use_linear_projection,\n \"class_embed_type\": class_embed_type,\n \"projection_class_embeddings_input_dim\": projection_class_embeddings_input_dim,\n }\n\n if not controlnet:\n config[\"out_channels\"] = unet_params.out_channels\n config[\"up_block_types\"] = tuple(up_block_types)\n\n return config\n\n\ndef create_vae_diffusers_config(original_config, image_size: int):\n \"\"\"\n Creates a config for the diffusers based on the config of the LDM model.\n \"\"\"\n vae_params = original_config.model.params.first_stage_config.params.ddconfig\n _ = original_config.model.params.first_stage_config.params.embed_dim\n\n block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult]\n down_block_types = [\"DownEncoderBlock2D\"] * len(block_out_channels)\n up_block_types = [\"UpDecoderBlock2D\"] * len(block_out_channels)\n\n config = {\n \"sample_size\": image_size,\n \"in_channels\": vae_params.in_channels,\n \"out_channels\": vae_params.out_ch,\n \"down_block_types\": tuple(down_block_types),\n \"up_block_types\": tuple(up_block_types),\n \"block_out_channels\": tuple(block_out_channels),\n \"latent_channels\": vae_params.z_channels,\n \"layers_per_block\": vae_params.num_res_blocks,\n }\n return config\n\n\ndef create_diffusers_schedular(original_config):\n schedular = DDIMScheduler(\n num_train_timesteps=original_config.model.params.timesteps,\n beta_start=original_config.model.params.linear_start,\n beta_end=original_config.model.params.linear_end,\n beta_schedule=\"scaled_linear\",\n )\n return schedular\n\n\ndef create_ldm_bert_config(original_config):\n bert_params = original_config.model.parms.cond_stage_config.params\n config = LDMBertConfig(\n d_model=bert_params.n_embed,\n encoder_layers=bert_params.n_layer,\n encoder_ffn_dim=bert_params.n_embed * 4,\n )\n return config\n\n\ndef convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False, controlnet=False):\n \"\"\"\n Takes a state dict and a config, and returns a converted checkpoint.\n \"\"\"\n\n # extract state_dict for UNet\n unet_state_dict = {}\n keys = list(checkpoint.keys())\n\n if controlnet:\n unet_key = \"control_model.\"\n else:\n unet_key = \"model.diffusion_model.\"\n\n # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA\n if sum(k.startswith(\"model_ema\") for k in keys) > 100 and extract_ema:\n print(f\"Checkpoint {path} has both EMA and non-EMA weights.\")\n print(\n \"In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA\"\n \" weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag.\"\n )\n for key in keys:\n if key.startswith(\"model.diffusion_model\"):\n flat_ema_key = \"model_ema.\" + \"\".join(key.split(\".\")[1:])\n unet_state_dict[key.replace(unet_key, \"\")] = checkpoint.pop(flat_ema_key)\n else:\n if sum(k.startswith(\"model_ema\") for k in keys) > 100:\n print(\n \"In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA\"\n \" weights (usually better for inference), please make sure to add the `--extract_ema` flag.\"\n )\n\n for key in keys:\n if key.startswith(unet_key):\n unet_state_dict[key.replace(unet_key, \"\")] = checkpoint.pop(key)\n\n new_checkpoint = {}\n\n new_checkpoint[\"time_embedding.linear_1.weight\"] = unet_state_dict[\"time_embed.0.weight\"]\n new_checkpoint[\"time_embedding.linear_1.bias\"] = unet_state_dict[\"time_embed.0.bias\"]\n new_checkpoint[\"time_embedding.linear_2.weight\"] = unet_state_dict[\"time_embed.2.weight\"]\n new_checkpoint[\"time_embedding.linear_2.bias\"] = unet_state_dict[\"time_embed.2.bias\"]\n\n if config[\"class_embed_type\"] is None:\n # No parameters to port\n ...\n elif config[\"class_embed_type\"] == \"timestep\" or config[\"class_embed_type\"] == \"projection\":\n new_checkpoint[\"class_embedding.linear_1.weight\"] = unet_state_dict[\"label_emb.0.0.weight\"]\n new_checkpoint[\"class_embedding.linear_1.bias\"] = unet_state_dict[\"label_emb.0.0.bias\"]\n new_checkpoint[\"class_embedding.linear_2.weight\"] = unet_state_dict[\"label_emb.0.2.weight\"]\n new_checkpoint[\"class_embedding.linear_2.bias\"] = unet_state_dict[\"label_emb.0.2.bias\"]\n else:\n raise NotImplementedError(f\"Not implemented `class_embed_type`: {config['class_embed_type']}\")\n\n new_checkpoint[\"conv_in.weight\"] = unet_state_dict[\"input_blocks.0.0.weight\"]\n new_checkpoint[\"conv_in.bias\"] = unet_state_dict[\"input_blocks.0.0.bias\"]\n\n if not controlnet:\n new_checkpoint[\"conv_norm_out.weight\"] = unet_state_dict[\"out.0.weight\"]\n new_checkpoint[\"conv_norm_out.bias\"] = unet_state_dict[\"out.0.bias\"]\n new_checkpoint[\"conv_out.weight\"] = unet_state_dict[\"out.2.weight\"]\n new_checkpoint[\"conv_out.bias\"] = unet_state_dict[\"out.2.bias\"]\n\n # Retrieves the keys for the input blocks only\n num_input_blocks = len({\".\".join(layer.split(\".\")[:2]) for layer in unet_state_dict if \"input_blocks\" in layer})\n input_blocks = {\n layer_id: [key for key in unet_state_dict if f\"input_blocks.{layer_id}\" in key]\n for layer_id in range(num_input_blocks)\n }\n\n # Retrieves the keys for the middle blocks only\n num_middle_blocks = len({\".\".join(layer.split(\".\")[:2]) for layer in unet_state_dict if \"middle_block\" in layer})\n middle_blocks = {\n layer_id: [key for key in unet_state_dict if f\"middle_block.{layer_id}\" in key]\n for layer_id in range(num_middle_blocks)\n }\n\n # Retrieves the keys for the output blocks only\n num_output_blocks = len({\".\".join(layer.split(\".\")[:2]) for layer in unet_state_dict if \"output_blocks\" in layer})\n output_blocks = {\n layer_id: [key for key in unet_state_dict if f\"output_blocks.{layer_id}\" in key]\n for layer_id in range(num_output_blocks)\n }\n\n for i in range(1, num_input_blocks):\n block_id = (i - 1) // (config[\"layers_per_block\"] + 1)\n layer_in_block_id = (i - 1) % (config[\"layers_per_block\"] + 1)\n\n resnets = [\n key for key in input_blocks[i] if f\"input_blocks.{i}.0\" in key and f\"input_blocks.{i}.0.op\" not in key\n ]\n attentions = [key for key in input_blocks[i] if f\"input_blocks.{i}.1\" in key]\n\n if f\"input_blocks.{i}.0.op.weight\" in unet_state_dict:\n new_checkpoint[f\"down_blocks.{block_id}.downsamplers.0.conv.weight\"] = unet_state_dict.pop(\n f\"input_blocks.{i}.0.op.weight\"\n )\n new_checkpoint[f\"down_blocks.{block_id}.downsamplers.0.conv.bias\"] = unet_state_dict.pop(\n f\"input_blocks.{i}.0.op.bias\"\n )\n\n paths = renew_resnet_paths(resnets)\n meta_path = {\"old\": f\"input_blocks.{i}.0\", \"new\": f\"down_blocks.{block_id}.resnets.{layer_in_block_id}\"}\n assign_to_checkpoint(\n paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config\n )\n\n if len(attentions):\n paths = renew_attention_paths(attentions)\n meta_path = {\"old\": f\"input_blocks.{i}.1\", \"new\": f\"down_blocks.{block_id}.attentions.{layer_in_block_id}\"}\n assign_to_checkpoint(\n paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config\n )\n\n resnet_0 = middle_blocks[0]\n attentions = middle_blocks[1]\n resnet_1 = middle_blocks[2]\n\n resnet_0_paths = renew_resnet_paths(resnet_0)\n assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config)\n\n resnet_1_paths = renew_resnet_paths(resnet_1)\n assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config)\n\n attentions_paths = renew_attention_paths(attentions)\n meta_path = {\"old\": \"middle_block.1\", \"new\": \"mid_block.attentions.0\"}\n assign_to_checkpoint(\n attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config\n )\n\n for i in range(num_output_blocks):\n block_id = i // (config[\"layers_per_block\"] + 1)\n layer_in_block_id = i % (config[\"layers_per_block\"] + 1)\n output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]]\n output_block_list = {}\n\n for layer in output_block_layers:\n layer_id, layer_name = layer.split(\".\")[0], shave_segments(layer, 1)\n if layer_id in output_block_list:\n output_block_list[layer_id].append(layer_name)\n else:\n output_block_list[layer_id] = [layer_name]\n\n if len(output_block_list) > 1:\n resnets = [key for key in output_blocks[i] if f\"output_blocks.{i}.0\" in key]\n attentions = [key for key in output_blocks[i] if f\"output_blocks.{i}.1\" in key]\n\n resnet_0_paths = renew_resnet_paths(resnets)\n paths = renew_resnet_paths(resnets)\n\n meta_path = {\"old\": f\"output_blocks.{i}.0\", \"new\": f\"up_blocks.{block_id}.resnets.{layer_in_block_id}\"}\n assign_to_checkpoint(\n paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config\n )\n\n output_block_list = {k: sorted(v) for k, v in output_block_list.items()}\n if [\"conv.bias\", \"conv.weight\"] in output_block_list.values():\n index = list(output_block_list.values()).index([\"conv.bias\", \"conv.weight\"])\n new_checkpoint[f\"up_blocks.{block_id}.upsamplers.0.conv.weight\"] = unet_state_dict[\n f\"output_blocks.{i}.{index}.conv.weight\"\n ]\n new_checkpoint[f\"up_blocks.{block_id}.upsamplers.0.conv.bias\"] = unet_state_dict[\n f\"output_blocks.{i}.{index}.conv.bias\"\n ]\n\n # Clear attentions as they have been attributed above.\n if len(attentions) == 2:\n attentions = []\n\n if len(attentions):\n paths = renew_attention_paths(attentions)\n meta_path = {\n \"old\": f\"output_blocks.{i}.1\",\n \"new\": f\"up_blocks.{block_id}.attentions.{layer_in_block_id}\",\n }\n assign_to_checkpoint(\n paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config\n )\n else:\n resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1)\n for path in resnet_0_paths:\n old_path = \".\".join([\"output_blocks\", str(i), path[\"old\"]])\n new_path = \".\".join([\"up_blocks\", str(block_id), \"resnets\", str(layer_in_block_id), path[\"new\"]])\n\n new_checkpoint[new_path] = unet_state_dict[old_path]\n\n if controlnet:\n # conditioning embedding\n\n orig_index = 0\n\n new_checkpoint[\"controlnet_cond_embedding.conv_in.weight\"] = unet_state_dict.pop(\n f\"input_hint_block.{orig_index}.weight\"\n )\n new_checkpoint[\"controlnet_cond_embedding.conv_in.bias\"] = unet_state_dict.pop(\n f\"input_hint_block.{orig_index}.bias\"\n )\n\n orig_index += 2\n\n diffusers_index = 0\n\n while diffusers_index < 6:\n new_checkpoint[f\"controlnet_cond_embedding.blocks.{diffusers_index}.weight\"] = unet_state_dict.pop(\n f\"input_hint_block.{orig_index}.weight\"\n )\n new_checkpoint[f\"controlnet_cond_embedding.blocks.{diffusers_index}.bias\"] = unet_state_dict.pop(\n f\"input_hint_block.{orig_index}.bias\"\n )\n diffusers_index += 1\n orig_index += 2\n\n new_checkpoint[\"controlnet_cond_embedding.conv_out.weight\"] = unet_state_dict.pop(\n f\"input_hint_block.{orig_index}.weight\"\n )\n new_checkpoint[\"controlnet_cond_embedding.conv_out.bias\"] = unet_state_dict.pop(\n f\"input_hint_block.{orig_index}.bias\"\n )\n\n # down blocks\n for i in range(num_input_blocks):\n new_checkpoint[f\"controlnet_down_blocks.{i}.weight\"] = unet_state_dict.pop(f\"zero_convs.{i}.0.weight\")\n new_checkpoint[f\"controlnet_down_blocks.{i}.bias\"] = unet_state_dict.pop(f\"zero_convs.{i}.0.bias\")\n\n # mid block\n new_checkpoint[\"controlnet_mid_block.weight\"] = unet_state_dict.pop(\"middle_block_out.0.weight\")\n new_checkpoint[\"controlnet_mid_block.bias\"] = unet_state_dict.pop(\"middle_block_out.0.bias\")\n\n return new_checkpoint\n\n\ndef convert_ldm_vae_checkpoint(checkpoint, config):\n # extract state dict for VAE\n vae_state_dict = {}\n vae_key = \"first_stage_model.\"\n keys = list(checkpoint.keys())\n for key in keys:\n if key.startswith(vae_key):\n vae_state_dict[key.replace(vae_key, \"\")] = checkpoint.get(key)\n\n new_checkpoint = {}\n\n new_checkpoint[\"encoder.conv_in.weight\"] = vae_state_dict[\"encoder.conv_in.weight\"]\n new_checkpoint[\"encoder.conv_in.bias\"] = vae_state_dict[\"encoder.conv_in.bias\"]\n new_checkpoint[\"encoder.conv_out.weight\"] = vae_state_dict[\"encoder.conv_out.weight\"]\n new_checkpoint[\"encoder.conv_out.bias\"] = vae_state_dict[\"encoder.conv_out.bias\"]\n new_checkpoint[\"encoder.conv_norm_out.weight\"] = vae_state_dict[\"encoder.norm_out.weight\"]\n new_checkpoint[\"encoder.conv_norm_out.bias\"] = vae_state_dict[\"encoder.norm_out.bias\"]\n\n new_checkpoint[\"decoder.conv_in.weight\"] = vae_state_dict[\"decoder.conv_in.weight\"]\n new_checkpoint[\"decoder.conv_in.bias\"] = vae_state_dict[\"decoder.conv_in.bias\"]\n new_checkpoint[\"decoder.conv_out.weight\"] = vae_state_dict[\"decoder.conv_out.weight\"]\n new_checkpoint[\"decoder.conv_out.bias\"] = vae_state_dict[\"decoder.conv_out.bias\"]\n new_checkpoint[\"decoder.conv_norm_out.weight\"] = vae_state_dict[\"decoder.norm_out.weight\"]\n new_checkpoint[\"decoder.conv_norm_out.bias\"] = vae_state_dict[\"decoder.norm_out.bias\"]\n\n new_checkpoint[\"quant_conv.weight\"] = vae_state_dict[\"quant_conv.weight\"]\n new_checkpoint[\"quant_conv.bias\"] = vae_state_dict[\"quant_conv.bias\"]\n new_checkpoint[\"post_quant_conv.weight\"] = vae_state_dict[\"post_quant_conv.weight\"]\n new_checkpoint[\"post_quant_conv.bias\"] = vae_state_dict[\"post_quant_conv.bias\"]\n\n # Retrieves the keys for the encoder down blocks only\n num_down_blocks = len({\".\".join(layer.split(\".\")[:3]) for layer in vae_state_dict if \"encoder.down\" in layer})\n down_blocks = {\n layer_id: [key for key in vae_state_dict if f\"down.{layer_id}\" in key] for layer_id in range(num_down_blocks)\n }\n\n # Retrieves the keys for the decoder up blocks only\n num_up_blocks = len({\".\".join(layer.split(\".\")[:3]) for layer in vae_state_dict if \"decoder.up\" in layer})\n up_blocks = {\n layer_id: [key for key in vae_state_dict if f\"up.{layer_id}\" in key] for layer_id in range(num_up_blocks)\n }\n\n for i in range(num_down_blocks):\n resnets = [key for key in down_blocks[i] if f\"down.{i}\" in key and f\"down.{i}.downsample\" not in key]\n\n if f\"encoder.down.{i}.downsample.conv.weight\" in vae_state_dict:\n new_checkpoint[f\"encoder.down_blocks.{i}.downsamplers.0.conv.weight\"] = vae_state_dict.pop(\n f\"encoder.down.{i}.downsample.conv.weight\"\n )\n new_checkpoint[f\"encoder.down_blocks.{i}.downsamplers.0.conv.bias\"] = vae_state_dict.pop(\n f\"encoder.down.{i}.downsample.conv.bias\"\n )\n\n paths = renew_vae_resnet_paths(resnets)\n meta_path = {\"old\": f\"down.{i}.block\", \"new\": f\"down_blocks.{i}.resnets\"}\n assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n\n mid_resnets = [key for key in vae_state_dict if \"encoder.mid.block\" in key]\n num_mid_res_blocks = 2\n for i in range(1, num_mid_res_blocks + 1):\n resnets = [key for key in mid_resnets if f\"encoder.mid.block_{i}\" in key]\n\n paths = renew_vae_resnet_paths(resnets)\n meta_path = {\"old\": f\"mid.block_{i}\", \"new\": f\"mid_block.resnets.{i - 1}\"}\n assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n\n mid_attentions = [key for key in vae_state_dict if \"encoder.mid.attn\" in key]\n paths = renew_vae_attention_paths(mid_attentions)\n meta_path = {\"old\": \"mid.attn_1\", \"new\": \"mid_block.attentions.0\"}\n assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n conv_attn_to_linear(new_checkpoint)\n\n for i in range(num_up_blocks):\n block_id = num_up_blocks - 1 - i\n resnets = [\n key for key in up_blocks[block_id] if f\"up.{block_id}\" in key and f\"up.{block_id}.upsample\" not in key\n ]\n\n if f\"decoder.up.{block_id}.upsample.conv.weight\" in vae_state_dict:\n new_checkpoint[f\"decoder.up_blocks.{i}.upsamplers.0.conv.weight\"] = vae_state_dict[\n f\"decoder.up.{block_id}.upsample.conv.weight\"\n ]\n new_checkpoint[f\"decoder.up_blocks.{i}.upsamplers.0.conv.bias\"] = vae_state_dict[\n f\"decoder.up.{block_id}.upsample.conv.bias\"\n ]\n\n paths = renew_vae_resnet_paths(resnets)\n meta_path = {\"old\": f\"up.{block_id}.block\", \"new\": f\"up_blocks.{i}.resnets\"}\n assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n\n mid_resnets = [key for key in vae_state_dict if \"decoder.mid.block\" in key]\n num_mid_res_blocks = 2\n for i in range(1, num_mid_res_blocks + 1):\n resnets = [key for key in mid_resnets if f\"decoder.mid.block_{i}\" in key]\n\n paths = renew_vae_resnet_paths(resnets)\n meta_path = {\"old\": f\"mid.block_{i}\", \"new\": f\"mid_block.resnets.{i - 1}\"}\n assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n\n mid_attentions = [key for key in vae_state_dict if \"decoder.mid.attn\" in key]\n paths = renew_vae_attention_paths(mid_attentions)\n meta_path = {\"old\": \"mid.attn_1\", \"new\": \"mid_block.attentions.0\"}\n assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n conv_attn_to_linear(new_checkpoint)\n return new_checkpoint\n\n\ndef convert_ldm_bert_checkpoint(checkpoint, config):\n def _copy_attn_layer(hf_attn_layer, pt_attn_layer):\n hf_attn_layer.q_proj.weight.data = pt_attn_layer.to_q.weight\n hf_attn_layer.k_proj.weight.data = pt_attn_layer.to_k.weight\n hf_attn_layer.v_proj.weight.data = pt_attn_layer.to_v.weight\n\n hf_attn_layer.out_proj.weight = pt_attn_layer.to_out.weight\n hf_attn_layer.out_proj.bias = pt_attn_layer.to_out.bias\n\n def _copy_linear(hf_linear, pt_linear):\n hf_linear.weight = pt_linear.weight\n hf_linear.bias = pt_linear.bias\n\n def _copy_layer(hf_layer, pt_layer):\n # copy layer norms\n _copy_linear(hf_layer.self_attn_layer_norm, pt_layer[0][0])\n _copy_linear(hf_layer.final_layer_norm, pt_layer[1][0])\n\n # copy attn\n _copy_attn_layer(hf_layer.self_attn, pt_layer[0][1])\n\n # copy MLP\n pt_mlp = pt_layer[1][1]\n _copy_linear(hf_layer.fc1, pt_mlp.net[0][0])\n _copy_linear(hf_layer.fc2, pt_mlp.net[2])\n\n def _copy_layers(hf_layers, pt_layers):\n for i, hf_layer in enumerate(hf_layers):\n if i != 0:\n i += i\n pt_layer = pt_layers[i : i + 2]\n _copy_layer(hf_layer, pt_layer)\n\n hf_model = LDMBertModel(config).eval()\n\n # copy embeds\n hf_model.model.embed_tokens.weight = checkpoint.transformer.token_emb.weight\n hf_model.model.embed_positions.weight.data = checkpoint.transformer.pos_emb.emb.weight\n\n # copy layer norm\n _copy_linear(hf_model.model.layer_norm, checkpoint.transformer.norm)\n\n # copy hidden layers\n _copy_layers(hf_model.model.layers, checkpoint.transformer.attn_layers.layers)\n\n _copy_linear(hf_model.to_logits, checkpoint.transformer.to_logits)\n\n return hf_model\n\n\ndef convert_ldm_clip_checkpoint_concise(checkpoint):\n keys = list(checkpoint.keys())\n text_model_dict = {}\n for key in keys:\n if key.startswith(\"cond_stage_model.transformer\"):\n text_model_dict[key[len(\"cond_stage_model.transformer.\") :]] = checkpoint[key]\n\n return text_model_dict\n\ndef convert_ldm_clip_checkpoint(checkpoint):\n text_model = CLIPTextModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n keys = list(checkpoint.keys())\n\n text_model_dict = {}\n\n for key in keys:\n if key.startswith(\"cond_stage_model.transformer\"):\n text_model_dict[key[len(\"cond_stage_model.transformer.\") :]] = checkpoint[key]\n\n text_model.load_state_dict(text_model_dict,strict=True)\n\n return text_model\n\n\ntextenc_conversion_lst = [\n (\"cond_stage_model.model.positional_embedding\", \"text_model.embeddings.position_embedding.weight\"),\n (\"cond_stage_model.model.token_embedding.weight\", \"text_model.embeddings.token_embedding.weight\"),\n (\"cond_stage_model.model.ln_final.weight\", \"text_model.final_layer_norm.weight\"),\n (\"cond_stage_model.model.ln_final.bias\", \"text_model.final_layer_norm.bias\"),\n]\ntextenc_conversion_map = {x[0]: x[1] for x in textenc_conversion_lst}\n\ntextenc_transformer_conversion_lst = [\n # (stable-diffusion, HF Diffusers)\n (\"resblocks.\", \"text_model.encoder.layers.\"),\n (\"ln_1\", \"layer_norm1\"),\n (\"ln_2\", \"layer_norm2\"),\n (\".c_fc.\", \".fc1.\"),\n (\".c_proj.\", \".fc2.\"),\n (\".attn\", \".self_attn\"),\n (\"ln_final.\", \"transformer.text_model.final_layer_norm.\"),\n (\"token_embedding.weight\", \"transformer.text_model.embeddings.token_embedding.weight\"),\n (\"positional_embedding\", \"transformer.text_model.embeddings.position_embedding.weight\"),\n]\nprotected = {re.escape(x[0]): x[1] for x in textenc_transformer_conversion_lst}\ntextenc_pattern = re.compile(\"|\".join(protected.keys()))\n\n\ndef convert_paint_by_example_checkpoint(checkpoint):\n config = CLIPVisionConfig.from_pretrained(\"openai/clip-vit-large-patch14\")\n model = PaintByExampleImageEncoder(config)\n\n keys = list(checkpoint.keys())\n\n text_model_dict = {}\n\n for key in keys:\n if key.startswith(\"cond_stage_model.transformer\"):\n text_model_dict[key[len(\"cond_stage_model.transformer.\") :]] = checkpoint[key]\n\n # load clip vision\n model.model.load_state_dict(text_model_dict)\n\n # load mapper\n keys_mapper = {\n k[len(\"cond_stage_model.mapper.res\") :]: v\n for k, v in checkpoint.items()\n if k.startswith(\"cond_stage_model.mapper\")\n }\n\n MAPPING = {\n \"attn.c_qkv\": [\"attn1.to_q\", \"attn1.to_k\", \"attn1.to_v\"],\n \"attn.c_proj\": [\"attn1.to_out.0\"],\n \"ln_1\": [\"norm1\"],\n \"ln_2\": [\"norm3\"],\n \"mlp.c_fc\": [\"ff.net.0.proj\"],\n \"mlp.c_proj\": [\"ff.net.2\"],\n }\n\n mapped_weights = {}\n for key, value in keys_mapper.items():\n prefix = key[: len(\"blocks.i\")]\n suffix = key.split(prefix)[-1].split(\".\")[-1]\n name = key.split(prefix)[-1].split(suffix)[0][1:-1]\n mapped_names = MAPPING[name]\n\n num_splits = len(mapped_names)\n for i, mapped_name in enumerate(mapped_names):\n new_name = \".\".join([prefix, mapped_name, suffix])\n shape = value.shape[0] // num_splits\n mapped_weights[new_name] = value[i * shape : (i + 1) * shape]\n\n model.mapper.load_state_dict(mapped_weights)\n\n # load final layer norm\n model.final_layer_norm.load_state_dict(\n {\n \"bias\": checkpoint[\"cond_stage_model.final_ln.bias\"],\n \"weight\": checkpoint[\"cond_stage_model.final_ln.weight\"],\n }\n )\n\n # load final proj\n model.proj_out.load_state_dict(\n {\n \"bias\": checkpoint[\"proj_out.bias\"],\n \"weight\": checkpoint[\"proj_out.weight\"],\n }\n )\n\n # load uncond vector\n model.uncond_vector.data = torch.nn.Parameter(checkpoint[\"learnable_vector\"])\n return model\n\n\ndef convert_open_clip_checkpoint(checkpoint):\n text_model = CLIPTextModel.from_pretrained(\"stabilityai/stable-diffusion-2\", subfolder=\"text_encoder\")\n\n keys = list(checkpoint.keys())\n\n text_model_dict = {}\n\n if \"cond_stage_model.model.text_projection\" in checkpoint:\n d_model = int(checkpoint[\"cond_stage_model.model.text_projection\"].shape[0])\n else:\n d_model = 1024\n\n text_model_dict[\"text_model.embeddings.position_ids\"] = text_model.text_model.embeddings.get_buffer(\"position_ids\")\n\n for key in keys:\n if \"resblocks.23\" in key: # Diffusers drops the final layer and only uses the penultimate layer\n continue\n if key in textenc_conversion_map:\n text_model_dict[textenc_conversion_map[key]] = checkpoint[key]\n if key.startswith(\"cond_stage_model.model.transformer.\"):\n new_key = key[len(\"cond_stage_model.model.transformer.\") :]\n if new_key.endswith(\".in_proj_weight\"):\n new_key = new_key[: -len(\".in_proj_weight\")]\n new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key)\n text_model_dict[new_key + \".q_proj.weight\"] = checkpoint[key][:d_model, :]\n text_model_dict[new_key + \".k_proj.weight\"] = checkpoint[key][d_model : d_model * 2, :]\n text_model_dict[new_key + \".v_proj.weight\"] = checkpoint[key][d_model * 2 :, :]\n elif new_key.endswith(\".in_proj_bias\"):\n new_key = new_key[: -len(\".in_proj_bias\")]\n new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key)\n text_model_dict[new_key + \".q_proj.bias\"] = checkpoint[key][:d_model]\n text_model_dict[new_key + \".k_proj.bias\"] = checkpoint[key][d_model : d_model * 2]\n text_model_dict[new_key + \".v_proj.bias\"] = checkpoint[key][d_model * 2 :]\n else:\n new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key)\n\n text_model_dict[new_key] = checkpoint[key]\n\n text_model.load_state_dict(text_model_dict)\n\n return text_model\n\n\ndef stable_unclip_image_encoder(original_config):\n \"\"\"\n Returns the image processor and clip image encoder for the img2img unclip pipeline.\n\n We currently know of two types of stable unclip models which separately use the clip and the openclip image\n encoders.\n \"\"\"\n\n image_embedder_config = original_config.model.params.embedder_config\n\n sd_clip_image_embedder_class = image_embedder_config.target\n sd_clip_image_embedder_class = sd_clip_image_embedder_class.split(\".\")[-1]\n\n if sd_clip_image_embedder_class == \"ClipImageEmbedder\":\n clip_model_name = image_embedder_config.params.model\n\n if clip_model_name == \"ViT-L/14\":\n feature_extractor = CLIPImageProcessor()\n image_encoder = CLIPVisionModelWithProjection.from_pretrained(\"openai/clip-vit-large-patch14\")\n else:\n raise NotImplementedError(f\"Unknown CLIP checkpoint name in stable diffusion checkpoint {clip_model_name}\")\n\n elif sd_clip_image_embedder_class == \"FrozenOpenCLIPImageEmbedder\":\n feature_extractor = CLIPImageProcessor()\n image_encoder = CLIPVisionModelWithProjection.from_pretrained(\"laion/CLIP-ViT-H-14-laion2B-s32B-b79K\")\n else:\n raise NotImplementedError(\n f\"Unknown CLIP image embedder class in stable diffusion checkpoint {sd_clip_image_embedder_class}\"\n )\n\n return feature_extractor, image_encoder\n\n\ndef stable_unclip_image_noising_components(\n original_config, clip_stats_path: Optional[str] = None, device: Optional[str] = None\n):\n \"\"\"\n Returns the noising components for the img2img and txt2img unclip pipelines.\n\n Converts the stability noise augmentor into\n 1. a `StableUnCLIPImageNormalizer` for holding the CLIP stats\n 2. a `DDPMScheduler` for holding the noise schedule\n\n If the noise augmentor config specifies a clip stats path, the `clip_stats_path` must be provided.\n \"\"\"\n noise_aug_config = original_config.model.params.noise_aug_config\n noise_aug_class = noise_aug_config.target\n noise_aug_class = noise_aug_class.split(\".\")[-1]\n\n if noise_aug_class == \"CLIPEmbeddingNoiseAugmentation\":\n noise_aug_config = noise_aug_config.params\n embedding_dim = noise_aug_config.timestep_dim\n max_noise_level = noise_aug_config.noise_schedule_config.timesteps\n beta_schedule = noise_aug_config.noise_schedule_config.beta_schedule\n\n image_normalizer = StableUnCLIPImageNormalizer(embedding_dim=embedding_dim)\n image_noising_scheduler = DDPMScheduler(num_train_timesteps=max_noise_level, beta_schedule=beta_schedule)\n\n if \"clip_stats_path\" in noise_aug_config:\n if clip_stats_path is None:\n raise ValueError(\"This stable unclip config requires a `clip_stats_path`\")\n\n clip_mean, clip_std = torch.load(clip_stats_path, map_location=device)\n clip_mean = clip_mean[None, :]\n clip_std = clip_std[None, :]\n\n clip_stats_state_dict = {\n \"mean\": clip_mean,\n \"std\": clip_std,\n }\n\n image_normalizer.load_state_dict(clip_stats_state_dict)\n else:\n raise NotImplementedError(f\"Unknown noise augmentor class: {noise_aug_class}\")\n\n return image_normalizer, image_noising_scheduler\n\n\ndef convert_controlnet_checkpoint(\n checkpoint, original_config, checkpoint_path, image_size, upcast_attention, extract_ema\n):\n ctrlnet_config = create_unet_diffusers_config(original_config, image_size=image_size, controlnet=True)\n ctrlnet_config[\"upcast_attention\"] = upcast_attention\n\n ctrlnet_config.pop(\"sample_size\")\n\n controlnet_model = ControlNetModel(**ctrlnet_config)\n\n converted_ctrl_checkpoint = convert_ldm_unet_checkpoint(\n checkpoint, ctrlnet_config, path=checkpoint_path, extract_ema=extract_ema, controlnet=True\n )\n\n controlnet_model.load_state_dict(converted_ctrl_checkpoint)\n\n return controlnet_model\n" }, { "path": "motionclone/utils/convert_lora_safetensor_to_diffusers.py", "content": "# coding=utf-8\n# Copyright 2023, Haofan Wang, Qixun Wang, All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# \n# Changes were made to this source code by Yuwei Guo.\n\"\"\" Conversion script for the LoRA's safetensors checkpoints. \"\"\"\n\nimport argparse\n\nimport torch\nfrom safetensors.torch import load_file\n\nfrom diffusers import StableDiffusionPipeline\n\n\ndef load_diffusers_lora(pipeline, state_dict, alpha=1.0):\n # directly update weight in diffusers model\n for key in state_dict:\n # only process lora down key\n if \"up.\" in key: continue\n\n up_key = key.replace(\".down.\", \".up.\")\n model_key = key.replace(\"processor.\", \"\").replace(\"_lora\", \"\").replace(\"down.\", \"\").replace(\"up.\", \"\")\n model_key = model_key.replace(\"to_out.\", \"to_out.0.\")\n layer_infos = model_key.split(\".\")[:-1]\n\n curr_layer = pipeline.unet\n while len(layer_infos) > 0:\n temp_name = layer_infos.pop(0)\n curr_layer = curr_layer.__getattr__(temp_name)\n\n weight_down = state_dict[key]\n weight_up = state_dict[up_key]\n curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).to(curr_layer.weight.data.device)\n\n return pipeline\n\n\ndef convert_lora(pipeline, state_dict, LORA_PREFIX_UNET=\"lora_unet\", LORA_PREFIX_TEXT_ENCODER=\"lora_te\", alpha=0.6):\n # load base model\n # pipeline = StableDiffusionPipeline.from_pretrained(base_model_path, torch_dtype=torch.float32)\n\n # load LoRA weight from .safetensors\n # state_dict = load_file(checkpoint_path)\n\n visited = []\n\n # directly update weight in diffusers model\n for key in state_dict:\n # it is suggested to print out the key, it usually will be something like below\n # \"lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight\"\n\n # as we have set the alpha beforehand, so just skip\n if \".alpha\" in key or key in visited:\n continue\n\n if \"text\" in key:\n layer_infos = key.split(\".\")[0].split(LORA_PREFIX_TEXT_ENCODER + \"_\")[-1].split(\"_\")\n curr_layer = pipeline.text_encoder\n else:\n layer_infos = key.split(\".\")[0].split(LORA_PREFIX_UNET + \"_\")[-1].split(\"_\")\n curr_layer = pipeline.unet\n\n # find the target layer\n temp_name = layer_infos.pop(0)\n while len(layer_infos) > -1:\n try:\n curr_layer = curr_layer.__getattr__(temp_name)\n if len(layer_infos) > 0:\n temp_name = layer_infos.pop(0)\n elif len(layer_infos) == 0:\n break\n except Exception:\n if len(temp_name) > 0:\n temp_name += \"_\" + layer_infos.pop(0)\n else:\n temp_name = layer_infos.pop(0)\n\n pair_keys = []\n if \"lora_down\" in key:\n pair_keys.append(key.replace(\"lora_down\", \"lora_up\"))\n pair_keys.append(key)\n else:\n pair_keys.append(key)\n pair_keys.append(key.replace(\"lora_up\", \"lora_down\"))\n\n # update weight\n if len(state_dict[pair_keys[0]].shape) == 4:\n weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)\n weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)\n curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3).to(curr_layer.weight.data.device)\n else:\n weight_up = state_dict[pair_keys[0]].to(torch.float32)\n weight_down = state_dict[pair_keys[1]].to(torch.float32)\n curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).to(curr_layer.weight.data.device)\n\n # update visited list\n for item in pair_keys:\n visited.append(item)\n\n return pipeline\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n\n parser.add_argument(\n \"--base_model_path\", default=None, type=str, required=True, help=\"Path to the base model in diffusers format.\"\n )\n parser.add_argument(\n \"--checkpoint_path\", default=None, type=str, required=True, help=\"Path to the checkpoint to convert.\"\n )\n parser.add_argument(\"--dump_path\", default=None, type=str, required=True, help=\"Path to the output model.\")\n parser.add_argument(\n \"--lora_prefix_unet\", default=\"lora_unet\", type=str, help=\"The prefix of UNet weight in safetensors\"\n )\n parser.add_argument(\n \"--lora_prefix_text_encoder\",\n default=\"lora_te\",\n type=str,\n help=\"The prefix of text encoder weight in safetensors\",\n )\n parser.add_argument(\"--alpha\", default=0.75, type=float, help=\"The merging ratio in W = W0 + alpha * deltaW\")\n parser.add_argument(\n \"--to_safetensors\", action=\"store_true\", help=\"Whether to store pipeline in safetensors format or not.\"\n )\n parser.add_argument(\"--device\", type=str, help=\"Device to use (e.g. cpu, cuda:0, cuda:1, etc.)\")\n\n args = parser.parse_args()\n\n base_model_path = args.base_model_path\n checkpoint_path = args.checkpoint_path\n dump_path = args.dump_path\n lora_prefix_unet = args.lora_prefix_unet\n lora_prefix_text_encoder = args.lora_prefix_text_encoder\n alpha = args.alpha\n\n pipe = convert(base_model_path, checkpoint_path, lora_prefix_unet, lora_prefix_text_encoder, alpha)\n\n pipe = pipe.to(args.device)\n pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)\n" }, { "path": "motionclone/utils/motionclone_functions.py", "content": "from dataclasses import dataclass\nimport os\nimport numpy as np\nimport torch\nimport matplotlib.pyplot as plt\nimport matplotlib.colors as mcolors\nfrom typing import Callable, List, Optional, Union\nfrom diffusers.utils import deprecate, logging, BaseOutput\nfrom .xformer_attention import *\nfrom .conv_layer import *\nfrom .util import *\nfrom diffusers.utils.torch_utils import randn_tensor\nfrom typing import List, Optional, Tuple, Union\nlogger = logging.get_logger(__name__) # pylint: disable=invalid-name\nfrom motionclone.utils.util import video_preprocess\nimport einops\nimport torchvision.transforms as transforms\n\ndef add_noise(self, timestep, x_0, noise_pred):\n alpha_prod_t = self.scheduler.alphas_cumprod[timestep]\n beta_prod_t = 1 - alpha_prod_t\n latents_noise = alpha_prod_t ** 0.5 * x_0 + beta_prod_t ** 0.5 * noise_pred\n return latents_noise\n \n@torch.no_grad()\ndef obtain_motion_representation(self, generator=None, motion_representation_path: str = None,\n duration=None,use_controlnet=False,):\n \n video_data = video_preprocess(self.input_config.video_path, self.input_config.height, \n self.input_config.width, self.input_config.video_length,duration=duration)\n video_latents = self.vae.encode(video_data.to(self.vae.dtype).to(self.vae.device)).latent_dist.mode()\n video_latents = self.vae.config.scaling_factor * video_latents\n video_latents = video_latents.unsqueeze(0)\n video_latents = einops.rearrange(video_latents, \"b f c h w -> b c f h w\")\n \n uncond_input = self.tokenizer(\n [\"\"], padding=\"max_length\", max_length=self.tokenizer.model_max_length,\n return_tensors=\"pt\"\n )\n step_t = int(self.input_config.add_noise_step)\n uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]\n noise_sampled = randn_tensor(video_latents.shape, generator=generator, device=video_latents.device, dtype=video_latents.dtype)\n noisy_latents = self.add_noise(step_t, video_latents, noise_sampled)\n \n down_block_additional_residuals = mid_block_additional_residual = None\n if use_controlnet:\n controlnet_image_index = self.input_config.image_index \n if self.controlnet.use_simplified_condition_embedding:\n controlnet_images = video_latents[:,:,controlnet_image_index,:,:] \n else:\n controlnet_images = (einops.rearrange(video_data.unsqueeze(0).to(self.vae.dtype).to(self.vae.device), \"b f c h w -> b c f h w\")+1)/2\n controlnet_images = controlnet_images[:,:,controlnet_image_index,:,:]\n\n controlnet_cond_shape = list(controlnet_images.shape)\n controlnet_cond_shape[2] = noisy_latents.shape[2]\n controlnet_cond = torch.zeros(controlnet_cond_shape).to(noisy_latents.device).to(noisy_latents.dtype)\n\n controlnet_conditioning_mask_shape = list(controlnet_cond.shape)\n controlnet_conditioning_mask_shape[1] = 1\n controlnet_conditioning_mask = torch.zeros(controlnet_conditioning_mask_shape).to(noisy_latents.device).to(noisy_latents.dtype)\n\n controlnet_cond[:,:,controlnet_image_index] = controlnet_images\n controlnet_conditioning_mask[:,:,controlnet_image_index] = 1\n\n down_block_additional_residuals, mid_block_additional_residual = self.controlnet(\n noisy_latents, step_t,\n encoder_hidden_states=uncond_embeddings,\n controlnet_cond=controlnet_cond,\n conditioning_mask=controlnet_conditioning_mask,\n conditioning_scale=self.input_config.controlnet_scale,\n guess_mode=False, return_dict=False,\n )\n\n _ = self.unet(noisy_latents, step_t, encoder_hidden_states=uncond_embeddings, return_dict=False, only_motion_feature=True,\n down_block_additional_residuals = down_block_additional_residuals,\n mid_block_additional_residual = mid_block_additional_residual,)\n temp_attn_prob_control = self.get_temp_attn_prob()\n \n motion_representation = { key: [max_value, max_index.to(torch.uint8)] for key, tensor in temp_attn_prob_control.items() for max_value, max_index in [torch.topk(tensor, k=1, dim=-1)]} \n \n torch.save(motion_representation, motion_representation_path)\n self.motion_representation_path = motion_representation_path\n\n\ndef compute_temp_loss(self, temp_attn_prob_control_dict):\n temp_attn_prob_loss = []\n for name in temp_attn_prob_control_dict.keys():\n current_temp_attn_prob = temp_attn_prob_control_dict[name]\n reference_representation_dict = self.motion_representation_dict[name]\n\n max_index = reference_representation_dict[1].to(torch.int64).to(current_temp_attn_prob.device)\n current_motion_representation = torch.gather(current_temp_attn_prob, index = max_index, dim=-1)\n \n reference_motion_representation = reference_representation_dict[0].to(dtype = current_motion_representation.dtype, device = current_motion_representation.device)\n \n module_attn_loss = F.mse_loss(current_motion_representation, reference_motion_representation.detach())\n temp_attn_prob_loss.append(module_attn_loss)\n \n loss_temp = torch.stack(temp_attn_prob_loss)\n return loss_temp.sum()\n \ndef sample_video(\n self,\n eta: float = 0.0,\n generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n noisy_latents: Optional[torch.FloatTensor] = None,\n add_controlnet: bool = False,\n):\n # Determine if use controlnet, i.e., conditional image2video\n self.add_controlnet = add_controlnet\n if self.add_controlnet:\n image_transforms = transforms.Compose([\n transforms.Resize((self.input_config.height, self.input_config.width)),\n transforms.ToTensor(),\n ])\n \n controlnet_images = [image_transforms(Image.open(path).convert(\"RGB\")) for path in self.input_config.condition_image_path_list]\n controlnet_images = torch.stack(controlnet_images).unsqueeze(0).to(dtype=self.vae.dtype,device=self.vae.device)\n controlnet_images = rearrange(controlnet_images, \"b f c h w -> b c f h w\")\n\n with torch.no_grad():\n if self.controlnet.use_simplified_condition_embedding:\n num_controlnet_images = controlnet_images.shape[2]\n controlnet_images = rearrange(controlnet_images, \"b c f h w -> (b f) c h w\")\n controlnet_images = self.vae.encode(controlnet_images * 2. - 1.).latent_dist.sample() * self.vae.config.scaling_factor\n self.controlnet_images = rearrange(controlnet_images, \"(b f) c h w -> b c f h w\", f=num_controlnet_images)\n else:\n self.controlnet_images = controlnet_images\n \n\n # Define call parameters\n # perform classifier_free_guidance in default\n batch_size = 1\n do_classifier_free_guidance = True\n device = self._execution_device\n \n # Encode input prompt\n self.text_embeddings = self._encode_prompt(self.input_config.new_prompt, device, 1, do_classifier_free_guidance, self.input_config.negative_prompt)\n # [uncond_embeddings, text_embeddings] [2, 77, 768]\n \n # Prepare latent variables\n noisy_latents = self.prepare_latents(\n batch_size,\n self.unet.config.in_channels,\n self.input_config.video_length,\n self.input_config.height,\n self.input_config.width,\n self.text_embeddings.dtype,\n device,\n generator,\n noisy_latents,\n )\n \n self.motion_representation_dict = torch.load(self.motion_representation_path)\n self.motion_scale = self.input_config.motion_guidance_weight\n\n extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)\n \n # save GPU memory\n # self.vae.to(device = \"cpu\")\n # self.text_encoder.to(device = \"cpu\")\n # torch.cuda.empty_cache()\n\n with self.progress_bar(total=self.input_config.inference_steps) as progress_bar:\n for step_index, step_t in enumerate(self.scheduler.timesteps):\n noisy_latents = self.single_step_video(noisy_latents, step_index, step_t, extra_step_kwargs) \n progress_bar.update()\n \n # decode latents for videos\n video = self.decode_latents(noisy_latents)\n return video\n\ndef single_step_video(self, noisy_latents, step_index, step_t, extra_step_kwargs):\n\n down_block_additional_residuals = mid_block_additional_residual = None\n if self.add_controlnet:\n with torch.no_grad(): \n controlnet_cond_shape = list(self.controlnet_images.shape)\n controlnet_cond_shape[2] = noisy_latents.shape[2]\n controlnet_cond = torch.zeros(controlnet_cond_shape).to(noisy_latents.device).to(noisy_latents.dtype)\n\n controlnet_conditioning_mask_shape = list(controlnet_cond.shape)\n controlnet_conditioning_mask_shape[1] = 1\n controlnet_conditioning_mask = torch.zeros(controlnet_conditioning_mask_shape).to(noisy_latents.device).to(noisy_latents.dtype)\n\n controlnet_image_index = self.input_config.image_index\n controlnet_cond[:,:,controlnet_image_index] = self.controlnet_images\n controlnet_conditioning_mask[:,:,controlnet_image_index] = 1\n\n down_block_additional_residuals, mid_block_additional_residual = self.controlnet(\n noisy_latents.expand(2,-1,-1,-1,-1), step_t,\n encoder_hidden_states=self.text_embeddings,\n controlnet_cond=controlnet_cond,\n conditioning_mask=controlnet_conditioning_mask,\n conditioning_scale=self.input_config.controlnet_scale,\n guess_mode=False, return_dict=False,\n )\n\n # Only require grad when need to compute the gradient for guidance\n if step_index < self.input_config.guidance_steps:\n \n down_block_additional_residuals_uncond = down_block_additional_residuals_cond = None\n mid_block_additional_residual_uncond = mid_block_additional_residual_cond = None\n if self.add_controlnet:\n down_block_additional_residuals_uncond = [tensor[[0],...].detach() for tensor in down_block_additional_residuals]\n down_block_additional_residuals_cond = [tensor[[1],...].detach() for tensor in down_block_additional_residuals]\n mid_block_additional_residual_uncond = mid_block_additional_residual[[0],...].detach()\n mid_block_additional_residual_cond = mid_block_additional_residual[[1],...].detach()\n \n control_latents = noisy_latents.clone().detach()\n control_latents.requires_grad = True\n\n control_latents = self.scheduler.scale_model_input(control_latents, step_t) \n noisy_latents = self.scheduler.scale_model_input(noisy_latents, step_t)\n \n with torch.no_grad():\n noise_pred_uncondition = self.unet(noisy_latents, step_t, encoder_hidden_states=self.text_embeddings[[0]],\n down_block_additional_residuals = down_block_additional_residuals_uncond,\n mid_block_additional_residual = mid_block_additional_residual_uncond,).sample.to(dtype=noisy_latents.dtype)\n\n noise_pred_condition = self.unet(control_latents, step_t, encoder_hidden_states=self.text_embeddings[[1]],\n down_block_additional_residuals = down_block_additional_residuals_cond,\n mid_block_additional_residual = mid_block_additional_residual_cond,).sample.to(dtype=noisy_latents.dtype)\n temp_attn_prob_control = self.get_temp_attn_prob()\n \n loss_motion = self.motion_scale * self.compute_temp_loss(temp_attn_prob_control,)\n \n if step_index < self.input_config.warm_up_steps:\n scale = (step_index+1)/self.input_config.warm_up_steps\n loss_motion = scale*loss_motion\n \n if step_index > self.input_config.guidance_steps-self.input_config.cool_up_steps:\n scale = (self.input_config.guidance_steps-step_index)/self.input_config.cool_up_steps\n loss_motion = scale*loss_motion\n\n gradient = torch.autograd.grad(loss_motion, control_latents, allow_unused=True)[0] # [1, 4, 16, 64, 64],\n assert gradient is not None, f\"Step {step_index}: grad is None\"\n\n noise_pred = noise_pred_condition + self.input_config.cfg_scale * (noise_pred_condition - noise_pred_uncondition)\n \n control_latents = self.scheduler.customized_step(noise_pred, step_index, control_latents, score=gradient.detach(),\n **extra_step_kwargs, return_dict=False)[0] # [1, 4, 16, 64, 64]\n return control_latents.detach()\n\n else:\n with torch.no_grad():\n noisy_latents = self.scheduler.scale_model_input(noisy_latents, step_t) \n noise_pred_group = self.unet(\n noisy_latents.expand(2,-1,-1,-1,-1), step_t, \n encoder_hidden_states=self.text_embeddings,\n down_block_additional_residuals = down_block_additional_residuals,\n mid_block_additional_residual = mid_block_additional_residual,\n ).sample.to(dtype=noisy_latents.dtype)\n\n noise_pred = noise_pred_group[[1]] + self.input_config.cfg_scale * (noise_pred_group[[1]] - noise_pred_group[[0]])\n noisy_latents = self.scheduler.customized_step(noise_pred, step_index, noisy_latents, score=None, **extra_step_kwargs, return_dict=False)[0] # [1, 4, 16, 64, 64]\n return noisy_latents.detach()\n \n\ndef get_temp_attn_prob(self,index_select=None):\n \n attn_prob_dic = {}\n\n for name, module in self.unet.named_modules():\n module_name = type(module).__name__\n if \"VersatileAttention\" in module_name and classify_blocks(self.input_config.motion_guidance_blocks, name):\n key = module.processor.key\n if index_select is not None:\n get_index = torch.repeat_interleave(torch.tensor(index_select), repeats=key.shape[0]//len(index_select))\n index_all = torch.arange(key.shape[0])\n index_picked = index_all[get_index.bool()]\n key = key[index_picked]\n key = module.reshape_heads_to_batch_dim(key).contiguous()\n \n query = module.processor.query\n if index_select is not None:\n query = query[index_picked]\n query = module.reshape_heads_to_batch_dim(query).contiguous()\n attention_probs = module.get_attention_scores(query, key, None) \n attention_probs = attention_probs.reshape(-1, module.heads,attention_probs.shape[1], attention_probs.shape[2]) \n attn_prob_dic[name] = attention_probs\n\n return attn_prob_dic\n\n@torch.no_grad()\ndef schedule_customized_step(\n self,\n model_output: torch.FloatTensor,\n step_index: int,\n sample: torch.FloatTensor,\n eta: float = 0.0,\n use_clipped_model_output: bool = False,\n generator=None,\n variance_noise: Optional[torch.FloatTensor] = None,\n return_dict: bool = True,\n\n # Guidance parameters\n score=None,\n guidance_scale=1.0,\n indices=None, # [0]\n return_middle = False,\n):\n if self.num_inference_steps is None:\n raise ValueError(\n \"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler\"\n )\n\n # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf\n # Ideally, read DDIM paper in-detail understanding\n\n # Notation ( -> \n # - pred_noise_t -> e_theta(x_t, t)\n # - pred_original_sample -> f_theta(x_t, t) or x_0\n # - std_dev_t -> sigma_t\n # - eta -> η\n # - pred_sample_direction -> \"direction pointing to x_t\"\n # - pred_prev_sample -> \"x_t-1\"\n\n \n # Support IF models\n if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in [\"learned\", \"learned_range\"]:\n model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)\n else:\n predicted_variance = None\n \n timestep = self.timesteps[step_index]\n # 1. get previous step value (=t-1)\n # prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps\n prev_timestep = self.timesteps[step_index+1] if step_index +1 = 0 else self.final_alpha_cumprod\n\n beta_prod_t = 1 - alpha_prod_t\n\n # 3. compute predicted original sample from predicted noise also called\n # \"predicted x_0\" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf\n if self.config.prediction_type == \"epsilon\":\n pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)\n pred_epsilon = model_output\n elif self.config.prediction_type == \"sample\":\n pred_original_sample = model_output\n pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)\n elif self.config.prediction_type == \"v_prediction\":\n pred_original_sample = (alpha_prod_t ** 0.5) * sample - (beta_prod_t ** 0.5) * model_output\n pred_epsilon = (alpha_prod_t ** 0.5) * model_output + (beta_prod_t ** 0.5) * sample\n else:\n raise ValueError(\n f\"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or\"\n \" `v_prediction`\"\n )\n\n # 4. Clip or threshold \"predicted x_0\"\n if self.config.thresholding:\n pred_original_sample = self._threshold_sample(pred_original_sample)\n elif self.config.clip_sample:\n pred_original_sample = pred_original_sample.clamp(\n -self.config.clip_sample_range, self.config.clip_sample_range\n )\n\n # 5. compute variance: \"sigma_t(η)\" -> see formula (16)\n # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)\n variance = self._get_variance(timestep, prev_timestep)\n std_dev_t = eta * variance ** (0.5)\n\n if use_clipped_model_output:\n # the pred_epsilon is always re-derived from the clipped x_0 in Glide\n pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) # [2, 4, 64, 64]\n \n if score is not None and return_middle:\n return pred_epsilon, alpha_prod_t, alpha_prod_t_prev, pred_original_sample\n\n # 6. apply guidance following the formula (14) from https://arxiv.org/pdf/2105.05233.pdf\n if score is not None and guidance_scale > 0.0: \n if indices is not None:\n # import pdb; pdb.set_trace()\n assert pred_epsilon[indices].shape == score.shape, \"pred_epsilon[indices].shape != score.shape\"\n pred_epsilon[indices] = pred_epsilon[indices] - guidance_scale * (1 - alpha_prod_t) ** (0.5) * score \n else:\n assert pred_epsilon.shape == score.shape\n pred_epsilon = pred_epsilon - guidance_scale * (1 - alpha_prod_t) ** (0.5) * score\n # \n\n # 7. compute \"direction pointing to x_t\" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf\n pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t ** 2) ** (0.5) * pred_epsilon \n\n # 8. compute x_t without \"random noise\" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf\n prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction \n\n if eta > 0:\n if variance_noise is not None and generator is not None:\n raise ValueError(\n \"Cannot pass both generator and variance_noise. Please make sure that either `generator` or\"\n \" `variance_noise` stays `None`.\"\n )\n\n if variance_noise is None:\n variance_noise = randn_tensor(\n model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype\n )\n variance = std_dev_t * variance_noise \n\n prev_sample = prev_sample + variance \n\n if not return_dict:\n return (prev_sample,)\n\n return prev_sample, pred_original_sample, alpha_prod_t_prev\n\n# equally to the function of schedule_customized_step (details refer to the discussion in https://github.com/LPengYang/MotionClone/issues/22)\n@torch.no_grad()\ndef schedule_customized_step_candidate(\n self,\n model_output: torch.FloatTensor,\n step_index: int,\n sample: torch.FloatTensor,\n eta: float = 0.0,\n use_clipped_model_output: bool = False,\n generator=None,\n variance_noise: Optional[torch.FloatTensor] = None,\n return_dict: bool = True,\n\n # Guidance parameters\n score=None,\n guidance_scale=1.0,\n indices=None, # [0]\n return_middle = False,\n):\n if self.num_inference_steps is None:\n raise ValueError(\n \"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler\"\n )\n\n # Support IF models\n if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in [\"learned\", \"learned_range\"]:\n model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)\n else:\n predicted_variance = None\n \n timestep = self.timesteps[step_index]\n # 1. get previous step value (=t-1)\n # prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps\n prev_timestep = self.timesteps[step_index+1] if step_index +1 = 0 else self.final_alpha_cumprod\n\n beta_prod_t = 1 - alpha_prod_t\n\n\n if self.config.prediction_type == \"epsilon\":\n if score is not None:\n alpha_t = alpha_prod_t/alpha_prod_t_prev\n model_output = model_output + guidance_scale/((1/alpha_t)**0.5-1)*score\n pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)\n pred_epsilon = model_output\n # print(\"validation\")\n \n else:\n pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)\n pred_epsilon = model_output\n elif self.config.prediction_type == \"sample\":\n pred_original_sample = model_output\n pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)\n elif self.config.prediction_type == \"v_prediction\":\n pred_original_sample = (alpha_prod_t ** 0.5) * sample - (beta_prod_t ** 0.5) * model_output\n pred_epsilon = (alpha_prod_t ** 0.5) * model_output + (beta_prod_t ** 0.5) * sample\n else:\n raise ValueError(\n f\"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or\"\n \" `v_prediction`\"\n )\n\n # 4. Clip or threshold \"predicted x_0\"\n if self.config.thresholding:\n pred_original_sample = self._threshold_sample(pred_original_sample)\n elif self.config.clip_sample:\n pred_original_sample = pred_original_sample.clamp(\n -self.config.clip_sample_range, self.config.clip_sample_range\n )\n\n # 5. compute variance: \"sigma_t(η)\" -> see formula (16)\n # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)\n variance = self._get_variance(timestep, prev_timestep)\n std_dev_t = eta * variance ** (0.5)\n\n if use_clipped_model_output:\n # the pred_epsilon is always re-derived from the clipped x_0 in Glide\n pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) # [2, 4, 64, 64]\n \n if score is not None and return_middle:\n return pred_epsilon, alpha_prod_t, alpha_prod_t_prev, pred_original_sample\n\n pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t ** 2) ** (0.5) * pred_epsilon \n\n prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction \n\n if eta > 0:\n if variance_noise is not None and generator is not None:\n raise ValueError(\n \"Cannot pass both generator and variance_noise. Please make sure that either `generator` or\"\n \" `variance_noise` stays `None`.\"\n )\n\n if variance_noise is None:\n variance_noise = randn_tensor(\n model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype\n )\n variance = std_dev_t * variance_noise \n\n prev_sample = prev_sample + variance \n\n if not return_dict:\n return (prev_sample,)\n\n return prev_sample, pred_original_sample, alpha_prod_t_prev\n\ndef schedule_set_timesteps(self, num_inference_steps: int, guidance_steps: int = 0, guiduance_scale: float = 0.0, device: Union[str, torch.device] = None,timestep_spacing_type= \"uneven\"):\n \"\"\"\n Sets the discrete timesteps used for the diffusion chain (to be run before inference).\n\n Args:\n num_inference_steps (`int`):\n The number of diffusion steps used when generating samples with a pre-trained model.\n \"\"\"\n\n if num_inference_steps > self.config.num_train_timesteps:\n raise ValueError(\n f\"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:\"\n f\" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle\"\n f\" maximal {self.config.num_train_timesteps} timesteps.\"\n )\n\n self.num_inference_steps = num_inference_steps\n \n # assign more steps in early denoising stage for motion guidance\n if timestep_spacing_type == \"uneven\":\n timesteps_guidance = (\n np.linspace(int((1-guiduance_scale)*self.config.num_train_timesteps), self.config.num_train_timesteps - 1, guidance_steps)\n .round()[::-1]\n .copy()\n .astype(np.int64)\n )\n timesteps_vanilla = (\n np.linspace(0, int((1-guiduance_scale)*self.config.num_train_timesteps) - 1, num_inference_steps-guidance_steps)\n .round()[::-1]\n .copy()\n .astype(np.int64)\n )\n timesteps = np.concatenate((timesteps_guidance, timesteps_vanilla))\n \n # \"linspace\", \"leading\", \"trailing\" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891\n elif timestep_spacing_type == \"linspace\":\n timesteps = (\n np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)\n .round()[::-1]\n .copy()\n .astype(np.int64)\n )\n elif timestep_spacing_type == \"leading\":\n step_ratio = self.config.num_train_timesteps // self.num_inference_steps\n # creates integer timesteps by multiplying by ratio\n # casting to int to avoid issues when num_inference_step is power of 3\n timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)\n timesteps += self.config.steps_offset\n elif timestep_spacing_type == \"trailing\":\n step_ratio = self.config.num_train_timesteps / self.num_inference_steps\n # creates integer timesteps by multiplying by ratio\n # casting to int to avoid issues when num_inference_step is power of 3\n timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)\n timesteps -= 1\n else:\n raise ValueError(\n f\"{timestep_spacing_type} is not supported. Please make sure to choose one of 'leading' or 'trailing'.\"\n )\n\n self.timesteps = torch.from_numpy(timesteps).to(device)\n \n@dataclass\nclass UNet3DConditionOutput(BaseOutput):\n sample: torch.FloatTensor\n\ndef unet_customized_forward(\n self,\n sample: torch.FloatTensor,\n timestep: Union[torch.Tensor, float, int],\n encoder_hidden_states: torch.Tensor,\n class_labels: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n\n # support controlnet\n down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,\n mid_block_additional_residual: Optional[torch.Tensor] = None,\n\n return_dict: bool = True,\n only_motion_feature: bool = False,\n ) -> Union[UNet3DConditionOutput, Tuple]:\n r\"\"\"\n Args:\n sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor\n timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps\n encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states\n return_dict (`bool`, *optional*, defaults to `True`):\n Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.\n\n Returns:\n [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:\n [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When\n returning a tuple, the first element is the sample tensor.\n \"\"\"\n # By default samples have to be AT least a multiple of the overall upsampling factor.\n # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).\n # However, the upsampling interpolation output size can be forced to fit any upsampling size\n # on the fly if necessary.\n default_overall_up_factor = 2**self.num_upsamplers\n\n # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`\n forward_upsample_size = False\n upsample_size = None\n\n if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):\n logger.info(\"Forward upsample size to force interpolation output size.\")\n forward_upsample_size = True\n\n # prepare attention_mask\n if attention_mask is not None:\n attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0\n attention_mask = attention_mask.unsqueeze(1)\n\n # center input if necessary\n if self.config.center_input_sample:\n sample = 2 * sample - 1.0\n\n # time\n timesteps = timestep\n if not torch.is_tensor(timesteps):\n # This would be a good case for the `match` statement (Python 3.10+)\n is_mps = sample.device.type == \"mps\"\n if isinstance(timestep, float):\n dtype = torch.float32 if is_mps else torch.float64\n else:\n dtype = torch.int32 if is_mps else torch.int64\n timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)\n elif len(timesteps.shape) == 0:\n timesteps = timesteps[None].to(sample.device)\n\n # broadcast to batch dimension in a way that's compatible with ONNX/Core ML\n timesteps = timesteps.expand(sample.shape[0])\n\n t_emb = self.time_proj(timesteps)\n\n # timesteps does not contain any weights and will always return f32 tensors\n # but time_embedding might actually be running in fp16. so we need to cast here.\n # there might be better ways to encapsulate this.\n t_emb = t_emb.to(dtype=self.dtype)\n emb = self.time_embedding(t_emb)\n\n if self.class_embedding is not None:\n if class_labels is None:\n raise ValueError(\"class_labels should be provided when num_class_embeds > 0\")\n\n if self.config.class_embed_type == \"timestep\":\n class_labels = self.time_proj(class_labels)\n\n class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)\n emb = emb + class_emb\n\n # pre-process\n sample = self.conv_in(sample)\n\n # down\n down_block_res_samples = (sample,)\n for downsample_block in self.down_blocks:\n if hasattr(downsample_block, \"has_cross_attention\") and downsample_block.has_cross_attention:\n sample, res_samples = downsample_block(\n hidden_states=sample,\n temb=emb,\n encoder_hidden_states=encoder_hidden_states,\n attention_mask=attention_mask,\n )\n else:\n sample, res_samples = downsample_block(hidden_states=sample, temb=emb, encoder_hidden_states=encoder_hidden_states)\n\n down_block_res_samples += res_samples\n\n # support controlnet\n down_block_res_samples = list(down_block_res_samples)\n if down_block_additional_residuals is not None:\n for i, down_block_additional_residual in enumerate(down_block_additional_residuals):\n if down_block_additional_residual.dim() == 4: # boardcast\n down_block_additional_residual = down_block_additional_residual.unsqueeze(2)\n down_block_res_samples[i] = down_block_res_samples[i] + down_block_additional_residual\n\n # mid\n sample = self.mid_block(\n sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask\n )\n\n # support controlnet\n if mid_block_additional_residual is not None:\n if mid_block_additional_residual.dim() == 4: # boardcast\n mid_block_additional_residual = mid_block_additional_residual.unsqueeze(2)\n sample = sample + mid_block_additional_residual\n \n # up\n for i, upsample_block in enumerate(self.up_blocks):\n if i<= int(self.input_config.motion_guidance_blocks[-1].split(\".\")[-1]):\n is_final_block = i == len(self.up_blocks) - 1\n\n res_samples = down_block_res_samples[-len(upsample_block.resnets) :]\n down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]\n\n # if we have not reached the final block and need to forward the\n # upsample size, we do it here\n if not is_final_block and forward_upsample_size:\n upsample_size = down_block_res_samples[-1].shape[2:]\n\n if hasattr(upsample_block, \"has_cross_attention\") and upsample_block.has_cross_attention:\n sample = upsample_block(\n hidden_states=sample,\n temb=emb,\n res_hidden_states_tuple=res_samples,\n encoder_hidden_states=encoder_hidden_states,\n upsample_size=upsample_size,\n attention_mask=attention_mask,\n )\n else:\n sample = upsample_block(\n hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size, encoder_hidden_states=encoder_hidden_states,\n )\n else:\n if only_motion_feature:\n return 0\n with torch.no_grad():\n is_final_block = i == len(self.up_blocks) - 1\n\n res_samples = down_block_res_samples[-len(upsample_block.resnets) :]\n down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]\n\n # if we have not reached the final block and need to forward the\n # upsample size, we do it here\n if not is_final_block and forward_upsample_size:\n upsample_size = down_block_res_samples[-1].shape[2:]\n\n if hasattr(upsample_block, \"has_cross_attention\") and upsample_block.has_cross_attention:\n sample = upsample_block(\n hidden_states=sample,\n temb=emb,\n res_hidden_states_tuple=res_samples,\n encoder_hidden_states=encoder_hidden_states,\n upsample_size=upsample_size,\n attention_mask=attention_mask,\n )\n else:\n sample = upsample_block(\n hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size, encoder_hidden_states=encoder_hidden_states,\n )\n \n # post-process\n sample = self.conv_norm_out(sample)\n sample = self.conv_act(sample)\n sample = self.conv_out(sample)\n\n if not return_dict:\n return (sample,)\n\n return UNet3DConditionOutput(sample=sample)\n\n" }, { "path": "motionclone/utils/util.py", "content": "import hashlib\nimport io\nimport re\nimport os\nimport imageio\nimport numpy as np\nfrom typing import Union\n\nimport cv2\nimport numpy as np\nimport requests\nimport random\nimport torch\nimport PIL.Image\nimport PIL.ImageOps\nfrom PIL import Image\nfrom typing import Callable, Union\n\nimport torch\nimport torchvision\nimport torch.distributed as dist\nimport torch.nn.functional as F\nimport decord\ndecord.bridge.set_bridge('torch')\nfrom PIL import Image, ImageOps\n\nfrom safetensors import safe_open\n# from tqdm import tqdm\nfrom einops import rearrange\nfrom motionclone.utils.convert_from_ckpt import convert_ldm_unet_checkpoint, convert_ldm_clip_checkpoint, convert_ldm_vae_checkpoint,convert_ldm_clip_checkpoint_concise\nfrom motionclone.utils.convert_lora_safetensor_to_diffusers import convert_lora, load_diffusers_lora\nfrom huggingface_hub import snapshot_download\n# from transformers import (\n# AutoFeatureExtractor,\n# BertTokenizerFast,\n# CLIPImageProcessor,\n# CLIPTextConfig,\n# CLIPTextModel,\n# CLIPTextModelWithProjection,\n# CLIPTokenizer,\n# CLIPVisionConfig,\n# CLIPVisionModelWithProjection,\n# )\n\nMOTION_MODULES = [\n \"mm_sd_v14.ckpt\", \n \"mm_sd_v15.ckpt\", \n \"mm_sd_v15_v2.ckpt\", \n \"v3_sd15_mm.ckpt\",\n]\n\nADAPTERS = [\n # \"mm_sd_v14.ckpt\",\n # \"mm_sd_v15.ckpt\",\n # \"mm_sd_v15_v2.ckpt\",\n # \"mm_sdxl_v10_beta.ckpt\",\n \"v2_lora_PanLeft.ckpt\",\n \"v2_lora_PanRight.ckpt\",\n \"v2_lora_RollingAnticlockwise.ckpt\",\n \"v2_lora_RollingClockwise.ckpt\",\n \"v2_lora_TiltDown.ckpt\",\n \"v2_lora_TiltUp.ckpt\",\n \"v2_lora_ZoomIn.ckpt\",\n \"v2_lora_ZoomOut.ckpt\",\n \"v3_sd15_adapter.ckpt\",\n # \"v3_sd15_mm.ckpt\",\n \"v3_sd15_sparsectrl_rgb.ckpt\",\n \"v3_sd15_sparsectrl_scribble.ckpt\",\n]\n\nBACKUP_DREAMBOOTH_MODELS = [\n \"realisticVisionV60B1_v51VAE.safetensors\",\n \"majicmixRealistic_v4.safetensors\",\n \"leosamsFilmgirlUltra_velvia20Lora.safetensors\",\n \"toonyou_beta3.safetensors\",\n \"majicmixRealistic_v5Preview.safetensors\",\n \"rcnzCartoon3d_v10.safetensors\",\n \"lyriel_v16.safetensors\",\n \"leosamsHelloworldXL_filmGrain20.safetensors\",\n \"TUSUN.safetensors\",\n]\n\ndef zero_rank_print(s):\n if (not dist.is_initialized()) and (dist.is_initialized() and dist.get_rank() == 0): print(\"### \" + s)\n\n\ndef save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8):\n videos = rearrange(videos, \"b c t h w -> t b c h w\")\n outputs = []\n for x in videos:\n x = torchvision.utils.make_grid(x, nrow=n_rows)\n x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)\n if rescale:\n x = (x + 1.0) / 2.0 # -1,1 -> 0,1\n x = (x * 255).numpy().astype(np.uint8)\n outputs.append(x)\n\n os.makedirs(os.path.dirname(path), exist_ok=True)\n imageio.mimsave(path, outputs, fps=fps)\n\ndef auto_download(local_path, is_dreambooth_lora=False):\n hf_repo = \"guoyww/animatediff_t2i_backups\" if is_dreambooth_lora else \"guoyww/animatediff\"\n folder, filename = os.path.split(local_path)\n\n if not os.path.exists(local_path):\n print(f\"local file {local_path} does not exist. trying to download from {hf_repo}\")\n\n if is_dreambooth_lora: assert filename in BACKUP_DREAMBOOTH_MODELS, f\"{filename} dose not exist in {hf_repo}\"\n else: assert filename in MOTION_MODULES + ADAPTERS, f\"{filename} dose not exist in {hf_repo}\"\n\n folder = \".\" if folder == \"\" else folder\n os.makedirs(folder, exist_ok=True)\n snapshot_download(repo_id=hf_repo, local_dir=folder, allow_patterns=[filename])\n\ndef load_weights(\n animation_pipeline,\n # motion module\n motion_module_path = \"\",\n motion_module_lora_configs = [],\n # domain adapter\n adapter_lora_path = \"\",\n adapter_lora_scale = 1.0,\n # image layers\n dreambooth_model_path = \"\",\n lora_model_path = \"\",\n lora_alpha = 0.8,\n):\n # motion module\n unet_state_dict = {}\n if motion_module_path != \"\":\n print(f\"load motion module from {motion_module_path}\")\n motion_module_state_dict = torch.load(motion_module_path, map_location=\"cpu\")\n motion_module_state_dict = motion_module_state_dict[\"state_dict\"] if \"state_dict\" in motion_module_state_dict else motion_module_state_dict\n unet_state_dict.update({name: param for name, param in motion_module_state_dict.items() if \"motion_modules.\" in name})\n unet_state_dict.pop(\"animatediff_config\", \"\")\n \n missing, unexpected = animation_pipeline.unet.load_state_dict(unet_state_dict, strict=False)\n # assert len(unexpected) == 0\n del unet_state_dict\n\n # base model\n if dreambooth_model_path != \"\":\n print(f\"load dreambooth model from {dreambooth_model_path}\")\n if dreambooth_model_path.endswith(\".safetensors\"):\n # import pdb; pdb.set_trace()\n dreambooth_state_dict = {}\n # import safetensors\n # dreambooth_state_dict = safetensors.torch.load_file(dreambooth_model_path)\n # import pdb; pdb.set_trace()\n with safe_open(dreambooth_model_path, framework=\"pt\", device=\"cpu\") as f:\n for key in f.keys():\n dreambooth_state_dict[key] = f.get_tensor(key)\n # import pdb; pdb.set_trace()\n elif dreambooth_model_path.endswith(\".ckpt\"):\n dreambooth_state_dict = torch.load(dreambooth_model_path, map_location=\"cpu\")\n \n # 1. vae\n converted_vae_checkpoint = convert_ldm_vae_checkpoint(dreambooth_state_dict, animation_pipeline.vae.config)\n animation_pipeline.vae.load_state_dict(converted_vae_checkpoint)\n # 2. unet\n converted_unet_checkpoint = convert_ldm_unet_checkpoint(dreambooth_state_dict, animation_pipeline.unet.config)\n animation_pipeline.unet.load_state_dict(converted_unet_checkpoint, strict=False)\n \n # 3. text_model\n # animation_pipeline.text_encoder = convert_ldm_clip_checkpoint(dreambooth_state_dict)\n converted_text_encoder_checkpoint = convert_ldm_clip_checkpoint_concise(dreambooth_state_dict)\n animation_pipeline.text_encoder.load_state_dict(converted_text_encoder_checkpoint, strict=True)\n del dreambooth_state_dict, converted_vae_checkpoint, converted_unet_checkpoint, converted_text_encoder_checkpoint\n \n # clip_config_name = \"models/clip-vit-large-patch14\"\n # clip_config = CLIPTextConfig.from_pretrained(clip_config_name, local_files_only=True)\n # text_model = CLIPTextModel(clip_config)\n # keys = list(dreambooth_state_dict.keys())\n # text_model_dict = {}\n # for key in keys:\n # if key.startswith(\"cond_stage_model.transformer\"):\n # text_model_dict[key[len(\"cond_stage_model.transformer.\") :]] = dreambooth_state_dict[key] \n # text_model.load_state_dict(text_model_dict)\n # animation_pipeline.text_encoder = text_model.to(dtype=animation_pipeline.unet.dtype)\n # # import pdb; pdb.set_trace()\n # # animation_pipeline.text_encoder = convert_ldm_clip_checkpoint(dreambooth_state_dict)\n # del dreambooth_state_dict\n \n # lora layers\n if lora_model_path != \"\":\n print(f\"load lora model from {lora_model_path}\")\n assert lora_model_path.endswith(\".safetensors\")\n lora_state_dict = {}\n with safe_open(lora_model_path, framework=\"pt\", device=\"cpu\") as f:\n for key in f.keys():\n lora_state_dict[key] = f.get_tensor(key)\n \n animation_pipeline = convert_lora(animation_pipeline, lora_state_dict, alpha=lora_alpha)\n del lora_state_dict\n\n # domain adapter lora\n if adapter_lora_path != \"\":\n print(f\"load domain lora from {adapter_lora_path}\")\n domain_lora_state_dict = torch.load(adapter_lora_path, map_location=\"cpu\")\n domain_lora_state_dict = domain_lora_state_dict[\"state_dict\"] if \"state_dict\" in domain_lora_state_dict else domain_lora_state_dict\n domain_lora_state_dict.pop(\"animatediff_config\", \"\")\n\n animation_pipeline = load_diffusers_lora(animation_pipeline, domain_lora_state_dict, alpha=adapter_lora_scale)\n\n # motion module lora\n for motion_module_lora_config in motion_module_lora_configs:\n path, alpha = motion_module_lora_config[\"path\"], motion_module_lora_config[\"alpha\"]\n print(f\"load motion LoRA from {path}\")\n motion_lora_state_dict = torch.load(path, map_location=\"cpu\")\n motion_lora_state_dict = motion_lora_state_dict[\"state_dict\"] if \"state_dict\" in motion_lora_state_dict else motion_lora_state_dict\n motion_lora_state_dict.pop(\"animatediff_config\", \"\")\n\n animation_pipeline = load_diffusers_lora(animation_pipeline, motion_lora_state_dict, alpha)\n\n return animation_pipeline\n\ndef video_preprocess(video_path, height, width, video_length, duration=None, sample_start_idx=0,):\n \n video_name = video_path.split('/')[-1].split('.')[0]\n vr = decord.VideoReader(video_path)\n fps = vr.get_avg_fps()\n if duration is None:\n # 读取整个视频\n total_frames = len(vr)\n else:\n # 根据给定的时长（秒）计算帧数\n total_frames = int(fps * duration)\n total_frames = min(total_frames, len(vr)) # 确保不超过视频总长度\n \n sample_index = np.linspace(0, total_frames - 1, video_length, dtype=int)\n print(total_frames,sample_index)\n video = vr.get_batch(sample_index)\n video = rearrange(video, \"f h w c -> f c h w\")\n\n video = F.interpolate(video, size=(height, width), mode=\"bilinear\", align_corners=True)\n \n # video_sample = rearrange(video, \"(b f) c h w -> b f h w c\", f=video_length)\n # imageio.mimwrite(f\"processed_videos/sample_{video_name}.mp4\", video_sample[0], fps=8, quality=9)\n \n video = video / 127.5 - 1.0\n\n return video\n\n\ndef set_nested_item(dataDict, mapList, value):\n \"\"\"Set item in nested dictionary\"\"\"\n \"\"\"\n Example: the mapList contains the name of each key ['injection','self-attn']\n this method will change the content in dataDict['injection']['self-attn'] with value\n\n \"\"\"\n for k in mapList[:-1]:\n dataDict = dataDict[k]\n dataDict[mapList[-1]] = value\n\n\ndef merge_sweep_config(base_config, update):\n \"\"\"Merge the updated parameters into the base config\"\"\"\n\n if base_config is None:\n raise ValueError(\"Base config is None\")\n if update is None:\n raise ValueError(\"Update config is None\")\n for key in update.keys():\n map_list = key.split(\"--\")\n set_nested_item(base_config, map_list, update[key])\n return base_config\n\n\n# Adapt from https://github.com/castorini/daam\ndef compute_token_merge_indices(tokenizer, prompt: str, word: str, word_idx: int = None, offset_idx: int = 0):\n merge_idxs = []\n tokens = tokenizer.tokenize(prompt.lower())\n if word_idx is None:\n word = word.lower()\n search_tokens = tokenizer.tokenize(word)\n start_indices = [x + offset_idx for x in range(len(tokens)) if\n tokens[x:x + len(search_tokens)] == search_tokens]\n for indice in start_indices:\n merge_idxs += [i + indice for i in range(0, len(search_tokens))]\n if not merge_idxs:\n raise Exception(f'Search word {word} not found in prompt!')\n else:\n merge_idxs.append(word_idx)\n\n return [x + 1 for x in merge_idxs], word_idx # Offset by 1.\n\n\ndef extract_data(input_string: str) -> list:\n print(\"input_string:\", input_string)\n \"\"\"\n Extract data from a string pattern where contents in () are separated by ;\n The first item in each () is considered as 'ref' and the rest as 'gen'.\n\n Args:\n - input_string (str): The input string pattern.\n\n Returns:\n - list: A list of dictionaries containing 'ref' and 'gen'.\n \"\"\"\n pattern = r'\\(([^)]+)\\)'\n matches = re.findall(pattern, input_string)\n\n data = []\n for match in matches:\n parts = [x.strip() for x in match.split(';')]\n ref = parts[0].strip()\n gen = parts[1].strip()\n data.append({'ref': ref, 'gen': gen})\n\n return data\n\n\ndef generate_hash_key(image, prompt=\"\"):\n \"\"\"\n Generate a hash key for the given image and prompt.\n \"\"\"\n byte_array = io.BytesIO()\n image.save(byte_array, format='JPEG')\n\n # Get byte data\n image_byte_data = byte_array.getvalue()\n\n # Combine image byte data and prompt byte data\n combined_data = image_byte_data + prompt.encode('utf-8')\n\n sha256 = hashlib.sha256()\n sha256.update(combined_data)\n return sha256.hexdigest()\n\n\ndef save_data(data, folder_path, key):\n \"\"\"\n Save data to a file, using key as the file name\n \"\"\"\n\n if not os.path.exists(folder_path):\n os.makedirs(folder_path)\n file_path = os.path.join(folder_path, f\"{key}.pt\")\n\n torch.save(data, file_path)\n\n\ndef get_data(folder_path, key):\n \"\"\"\n Get data from a file, using key as the file name\n :param folder_path:\n :param key:\n :return:\n \"\"\"\n\n file_path = os.path.join(folder_path, f\"{key}.pt\")\n if os.path.exists(file_path):\n return torch.load(file_path)\n else:\n return None\n\n\ndef PILtoTensor(data: Image.Image) -> torch.Tensor:\n return torch.tensor(np.array(data)).permute(2, 0, 1).unsqueeze(0).float()\n\n\ndef TensorToPIL(data: torch.Tensor) -> Image.Image:\n return Image.fromarray(data.squeeze().permute(1, 2, 0).numpy().astype(np.uint8))\n\n# Adapt from https://github.com/huggingface/diffusers/blob/v0.26.3/src/diffusers/utils/loading_utils.py#L9\ndef load_image(\n image: Union[str, PIL.Image.Image], convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] = None\n) -> PIL.Image.Image:\n \"\"\"\n Loads `image` to a PIL Image.\n\n Args:\n image (`str` or `PIL.Image.Image`):\n The image to convert to the PIL Image format.\n convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], optional):\n A conversion method to apply to the image after loading it.\n When set to `None` the image will be converted \"RGB\".\n\n Returns:\n `PIL.Image.Image`:\n A PIL Image.\n \"\"\"\n if isinstance(image, str):\n if image.startswith(\"http://\") or image.startswith(\"https://\"):\n image = PIL.Image.open(requests.get(image, stream=True).raw)\n elif os.path.isfile(image):\n image = PIL.Image.open(image)\n else:\n raise ValueError(\n f\"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path.\"\n )\n elif isinstance(image, PIL.Image.Image):\n image = image\n else:\n raise ValueError(\n \"Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image.\"\n )\n\n image = PIL.ImageOps.exif_transpose(image)\n\n if convert_method is not None:\n image = convert_method(image)\n else:\n image = image.convert(\"RGB\")\n\n return image\n\n# Take from huggingface/diffusers\ndef rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):\n \"\"\"\n Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and\n Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4\n \"\"\"\n std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)\n std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)\n # rescale the results from guidance (fixes overexposure)\n noise_pred_rescaled = noise_cfg * (std_text / std_cfg)\n # mix with the original results from guidance by factor guidance_rescale to avoid \"plain looking\" images\n noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg\n return noise_cfg\n\ndef _in_step(config, step):\n in_step = False\n try:\n start_step = config.start_step\n end_step = config.end_step\n if start_step <= step < end_step:\n in_step = True\n except:\n in_step = False\n return in_step\n\ndef classify_blocks(block_list, name):\n is_correct_block = False\n for block in block_list:\n if block in name:\n is_correct_block = True\n break\n return is_correct_block\n\ndef set_all_seed(seed):\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n np.random.seed(seed)\n random.seed(seed)\n torch.backends.cudnn.deterministic = True" }, { "path": "motionclone/utils/xformer_attention.py", "content": "import math\nfrom typing import Optional, Callable\nimport xformers\nfrom omegaconf import OmegaConf\nimport yaml\nfrom .util import classify_blocks\n\ndef identify_blocks(block_list, name):\n block_name = None\n for block in block_list:\n if block in name:\n block_name = block\n break\n return block_name\n\n\nclass MySelfAttnProcessor:\n def __init__(self, attention_op: Optional[Callable] = None):\n self.attention_op = attention_op\n \n\n def __call__(self, attn, hidden_states, query, key, value, attention_mask):\n # self.attn = attn\n self.key = key\n self.query = query\n # self.value = value\n # self.attention_mask = attention_mask\n # self.hidden_state = hidden_states.detach()\n # return hidden_states\n \n def record_qkv(self, attn, hidden_states, query, key, value, attention_mask):\n # self.attn = attn\n self.key = key\n self.query = query\n # self.value = value\n # # self.attention_mask = attention_mask\n # self.hidden_state = hidden_states.detach()\n # # import pdb; pdb.set_trace()\n \n def record_attn_mask(self, attn, hidden_states, query, key, value, attention_mask):\n self.attn = attn\n self.attention_mask = attention_mask\n \n\ndef prep_unet_attention(unet,motion_gudiance_blocks):\n # replace the fwd function\n for name, module in unet.named_modules():\n module_name = type(module).__name__\n if \"VersatileAttention\" in module_name and classify_blocks(motion_gudiance_blocks, name): # the temporary attention in guidance blocks\n module.set_processor(MySelfAttnProcessor())\n # print(module_name)\n return unet\n\n\ndef get_self_attn_feat(unet, injection_config, config):\n hidden_state_dict = dict()\n query_dict = dict()\n key_dict = dict()\n value_dict = dict()\n \n for name, module in unet.named_modules():\n module_name = type(module).__name__\n if \"CrossAttention\" in module_name and 'attn1' in name and classify_blocks(injection_config.blocks, name=name):\n res = int(math.sqrt(module.processor.hidden_state.shape[1]))\n # import pdb; pdb.set_trace()\n bs = module.processor.hidden_state.shape[0] # 20 * 16 = 320\n # block_name = identify_blocks(injection_config.blocks, name=name)\n # block_id = int(block_name.split('.')[-1])\n # h = config.H // (32 * block_id)\n # w = config.W // (32 * block_id)\n hidden_state_dict[name] = module.processor.hidden_state.cpu().permute(0, 2, 1).reshape(bs, -1, res, res)\n res = int(math.sqrt(module.processor.query.shape[1]))\n query_dict[name] = module.processor.query.cpu().permute(0, 2, 1).reshape(bs, -1, res, res)\n key_dict[name] = module.processor.key.cpu().permute(0, 2, 1).reshape(bs, -1, res, res)\n value_dict[name] = module.processor.value.cpu().permute(0, 2, 1).reshape(bs, -1, res, res)\n # import pdb; pdb.set_trace()\n # import pdb; pdb.set_trace()\n return hidden_state_dict, query_dict, key_dict, value_dict\n\n\ndef clean_attn_buffer(unet):\n for name, module in unet.named_modules():\n module_name = type(module).__name__\n if module_name == \"Attention\" and 'attn' in name:\n if 'injection_config' in module.processor.__dict__.keys():\n module.processor.injection_config = None\n if 'injection_mask' in module.processor.__dict__.keys():\n module.processor.injection_mask = None\n if 'obj_index' in module.processor.__dict__.keys():\n module.processor.obj_index = None\n if 'pca_weight' in module.processor.__dict__.keys():\n module.processor.pca_weight = None\n if 'pca_weight_changed' in module.processor.__dict__.keys():\n module.processor.pca_weight_changed = None\n if 'pca_info' in module.processor.__dict__.keys():\n module.processor.pca_info = None\n if 'step' in module.processor.__dict__.keys():\n module.processor.step = None\n" }, { "path": "t2v_video_sample.py", "content": "import argparse\nfrom omegaconf import OmegaConf\nimport torch\nfrom diffusers import AutoencoderKL, DDIMScheduler\nfrom transformers import CLIPTextModel, CLIPTokenizer\nfrom motionclone.models.unet import UNet3DConditionModel\nfrom motionclone.pipelines.pipeline_animation import AnimationPipeline\nfrom motionclone.utils.util import load_weights\nfrom diffusers.utils.import_utils import is_xformers_available\nfrom motionclone.utils.motionclone_functions import *\nimport json\nfrom motionclone.utils.xformer_attention import *\n\ndef main(args):\n\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = args.visible_gpu or str(os.getenv('CUDA_VISIBLE_DEVICES', 0))\n \n config = OmegaConf.load(args.inference_config)\n adopted_dtype = torch.float16\n device = \"cuda\"\n set_all_seed(42)\n\n tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_path, subfolder=\"tokenizer\")\n text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_path, subfolder=\"text_encoder\").to(device).to(dtype=adopted_dtype)\n vae = AutoencoderKL.from_pretrained(args.pretrained_model_path, subfolder=\"vae\").to(device).to(dtype=adopted_dtype)\n \n config.width = config.get(\"W\", args.W)\n config.height = config.get(\"H\", args.H)\n config.video_length = config.get(\"L\", args.L)\n \n if not os.path.exists(args.generated_videos_save_dir):\n os.makedirs(args.generated_videos_save_dir)\n OmegaConf.save(config, os.path.join(args.generated_videos_save_dir,\"inference_config.json\"))\n \n model_config = OmegaConf.load(config.get(\"model_config\", \"\"))\n unet = UNet3DConditionModel.from_pretrained_2d(args.pretrained_model_path, subfolder=\"unet\", unet_additional_kwargs=OmegaConf.to_container(model_config.unet_additional_kwargs),).to(device).to(dtype=adopted_dtype)\n \n # set xformers\n if is_xformers_available() and (not args.without_xformers):\n unet.enable_xformers_memory_efficient_attention()\n\n pipeline = AnimationPipeline(\n vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet,\n controlnet=None,\n scheduler=DDIMScheduler(**OmegaConf.to_container(model_config.noise_scheduler_kwargs)),\n ).to(device)\n \n pipeline = load_weights(\n pipeline,\n # motion module\n motion_module_path = config.get(\"motion_module\", \"\"),\n dreambooth_model_path = config.get(\"dreambooth_path\", \"\"),\n ).to(device)\n pipeline.text_encoder.to(dtype=adopted_dtype)\n \n # load customized functions from motionclone_functions\n pipeline.scheduler.customized_step = schedule_customized_step.__get__(pipeline.scheduler)\n pipeline.scheduler.customized_set_timesteps = schedule_set_timesteps.__get__(pipeline.scheduler)\n pipeline.unet.forward = unet_customized_forward.__get__(pipeline.unet)\n pipeline.sample_video = sample_video.__get__(pipeline)\n pipeline.single_step_video = single_step_video.__get__(pipeline)\n pipeline.get_temp_attn_prob = get_temp_attn_prob.__get__(pipeline)\n pipeline.add_noise = add_noise.__get__(pipeline)\n pipeline.compute_temp_loss = compute_temp_loss.__get__(pipeline)\n pipeline.obtain_motion_representation = obtain_motion_representation.__get__(pipeline)\n \n for param in pipeline.unet.parameters():\n param.requires_grad = False\n pipeline.input_config, pipeline.unet.input_config = config, config\n \n pipeline.unet = prep_unet_attention(pipeline.unet,pipeline.input_config.motion_guidance_blocks)\n pipeline.unet = prep_unet_conv(pipeline.unet)\n pipeline.scheduler.customized_set_timesteps(config.inference_steps, config.guidance_steps,config.guidance_scale,device=device,timestep_spacing_type = \"uneven\")\n # pipeline.scheduler.customized_set_timesteps(config.inference_steps,device=device,timestep_spacing_type = \"linspace\")\n with open(args.examples, 'r') as files:\n for line in files:\n # prepare infor of each case\n example_infor = json.loads(line)\n config.video_path = example_infor[\"video_path\"]\n config.new_prompt = example_infor[\"new_prompt\"] + config.get(\"positive_prompt\", \"\")\n pipeline.input_config, pipeline.unet.input_config = config, config # update config\n \n # perform motion representation extraction\n seed_motion = example_infor.get(\"seed\", args.default_seed) \n generator = torch.Generator(device=pipeline.device)\n generator.manual_seed(seed_motion)\n if not os.path.exists(args.motion_representation_save_dir):\n os.makedirs(args.motion_representation_save_dir)\n motion_representation_path = os.path.join(args.motion_representation_save_dir, os.path.splitext(os.path.basename(config.video_path))[0] + '.pt') \n pipeline.obtain_motion_representation(generator= generator, motion_representation_path = motion_representation_path) \n \n # perform video generation \n seed = seed_motion # can assign other seed here\n generator = torch.Generator(device=pipeline.device)\n generator.manual_seed(seed)\n pipeline.input_config.seed = seed\n \n videos = pipeline.sample_video(generator = generator,)\n videos = rearrange(videos, \"b c f h w -> b f h w c\")\n save_path = os.path.join(args.generated_videos_save_dir, os.path.splitext(os.path.basename(config.video_path))[0] \n + \"_\" + config.new_prompt.strip().replace(' ', '_') + str(seed_motion) + \"_\" +str(seed)+'.mp4') \n videos_uint8 = (videos[0] * 255).astype(np.uint8)\n\n imageio.mimwrite(save_path, videos_uint8, fps=8)\n print(save_path,\"is done\")\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--pretrained-model-path\", type=str, default=\"models/StableDiffusion\",)\n \n parser.add_argument(\"--inference_config\", type=str, default=\"configs/t2v_camera.yaml\")\n parser.add_argument(\"--examples\", type=str, default=\"configs/t2v_camera.jsonl\")\n parser.add_argument(\"--motion-representation-save-dir\", type=str, default=\"motion_representation/\")\n parser.add_argument(\"--generated-videos-save-dir\", type=str, default=\"generated_videos\")\n \n parser.add_argument(\"--visible_gpu\", type=str, default=None)\n parser.add_argument(\"--default-seed\", type=int, default=2025)\n parser.add_argument(\"--L\", type=int, default=16)\n parser.add_argument(\"--W\", type=int, default=512)\n parser.add_argument(\"--H\", type=int, default=512)\n\n parser.add_argument(\"--without-xformers\", action=\"store_true\")\n\n args = parser.parse_args()\n main(args)\n" } ]