Version 0.1.0 Release

.github/dependabot.yml

@@ -0,0 +1,8 @@
+version: 2
+updates:
+  - package-ecosystem: github-actions
+    directory: /
+    schedule:
+      interval: daily
+    commit-message:
+      prefix: chore

.github/workflows/test.yaml

@@ -0,0 +1,24 @@
+name: test
+on:
+  workflow_dispatch:
+  pull_request:
+    branches: ['*', '!latest']
+  push:
+    branches: [latest, main]
+jobs:
+  lib:
+    runs-on: [self-hosted, linux, x64]
+    steps:
+      - name: Checkout
+        uses: actions/[email protected]
+        with:
+          submodules: true
+
+      - name: Run tests
+        shell: fish {0}
+        run: |
+          conda activate minimal-animatediff
+          make test
+
+      - name: Check diff
+        run: git diff --exit-code HEAD

.gitignore

@@ -1,2 +1,10 @@
 .*
 !.git*
+!.pylintrc
+
+__pycache__/
+coverage/
+flagged/
+models/
+samples/
+snapshots/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "deps/AnimateDiff"]
+	path = deps/AnimateDiff
+	url = https://github.com/guoyww/AnimateDiff

Makefile

@@ -0,0 +1,17 @@
+format:
+	python3 -m black minimal_animatediff scripts tests \
+		--line-length 120
+
+lint: format
+	python3 -m flake8 minimal_animatediff scripts tests \
+		--max-line-length 120
+
+test: lint
+	python3 -m pytest  -v -s \
+		--durations=0 \
+		--cov=minimal_animatediff \
+		--cov-fail-under=85 \
+		--cov-report term \
+		--cov-report annotate:coverage
+
+.PHONY: format lint test

environment.yaml

@@ -0,0 +1,26 @@
+name: minimal-animatediff
+channels:
+  - pytorch
+  - xformers
+dependencies:
+  - cudatoolkit=11.3
+  - pip
+  - python=3.10
+  - pytorch==1.12.1
+  - torchvision==0.13.1
+  - xformers
+  - pip:
+    - black
+    - diffusers[torch]==0.11.1
+    - einops
+    - filehash
+    - flake8
+    - gdown
+    - huggingface_hub
+    - imageio[ffmpeg]
+    - pytest
+    - pytest-cov
+    - pytest-dependency
+    - safetensors
+    - transformers==4.25.1
+    - triton

minimal_animatediff/animation_pipeline.py

@@ -0,0 +1,54 @@
+import sys
+
+from diffusers import DDIMScheduler
+
+from deps.AnimateDiff.animatediff.pipelines.pipeline_animation import AnimationPipeline
+import deps.AnimateDiff.animatediff.utils.convert_from_ckpt as cvt
+
+from .dream_booth import DreamBoothModel
+from .motion_module import MotionModule
+from .stable_diffusion import StableDiffusion
+
+
+def create_animation_pipeline():
+    sd = StableDiffusion()
+
+    pipeline = AnimationPipeline(
+        text_encoder=sd.text_encoder,
+        tokenizer=sd.tokenizer,
+        vae=sd.vae,
+        unet=sd.unet,
+        scheduler=DDIMScheduler(
+            **{
+                "num_train_timesteps": 1000,
+                "beta_start": 0.00085,
+                "beta_end": 0.012,
+                "beta_schedule": "linear",
+                "steps_offset": 1,
+                "clip_sample": False,
+            }
+        ),
+    )
+
+    pipeline.to("cuda")
+
+    print("Loading motion module to the animation pipeline...")
+    mm = MotionModule("mm_sd_v15.ckpt")
+    _, unexpected = pipeline.unet.load_state_dict(mm.states, strict=False)
+    if len(unexpected) > 0:
+        sys.exit("Failed to load motion module to the animation pipeline!")
+
+    print("Loading diffusion model to the animation pipeline...")
+    db_model = DreamBoothModel("toonyou_beta3.safetensors")
+
+    converted_vae_checkpoint = cvt.convert_ldm_vae_checkpoint(db_model.states, pipeline.vae.config)
+    pipeline.vae.load_state_dict(converted_vae_checkpoint)
+
+    converted_unet_checkpoint = cvt.convert_ldm_unet_checkpoint(db_model.states, pipeline.unet.config)
+    pipeline.unet.load_state_dict(converted_unet_checkpoint, strict=False)
+
+    pipeline.text_encoder = cvt.convert_ldm_clip_checkpoint(db_model.states)
+
+    pipeline.to("cuda")
+
+    return pipeline

minimal_animatediff/dream_booth.py

@@ -0,0 +1,13 @@
+from safetensors import safe_open
+from .utils import populate_snapshot
+
+
+class DreamBoothModel:
+    def __init__(self, name: str):
+        snapshot = populate_snapshot("dream_booth/" + name)
+
+        # Load the model states
+        self.states = {}
+        with safe_open(snapshot, framework="pt", device="cpu") as model:
+            for key in model.keys():
+                self.states[key] = model.get_tensor(key)

minimal_animatediff/motion_module.py

@@ -0,0 +1,9 @@
+import torch
+
+from .utils import populate_snapshot
+
+
+class MotionModule:
+    def __init__(self, name: str):
+        snapshot = populate_snapshot("motion_modules/" + name)
+        self.states = torch.load(snapshot, map_location="cpu")

minimal_animatediff/stable_diffusion.py

@@ -0,0 +1,53 @@
+from diffusers import AutoencoderKL
+from diffusers.utils.import_utils import is_xformers_available
+from huggingface_hub import snapshot_download
+from transformers import CLIPTextModel, CLIPTokenizer
+
+from deps.AnimateDiff.animatediff.models.unet import UNet3DConditionModel
+
+
+class StableDiffusion:
+    def __init__(self):
+        snapshot = "snapshots/stable_diffusion"
+        snapshot_download(
+            repo_id="runwayml/stable-diffusion-v1-5",
+            local_dir=snapshot,
+            allow_patterns=[
+                "text_encoder/*.json",
+                "text_encoder/*model.bin",
+                "tokenizer/*",
+                "unet/*.json",
+                "unet/*model.bin",
+                "vae/*.json",
+                "vae/*model.bin",
+            ],
+        )
+
+        self.text_encoder = CLIPTextModel.from_pretrained(snapshot, subfolder="text_encoder")
+        self.tokenizer = CLIPTokenizer.from_pretrained(snapshot, subfolder="tokenizer")
+        self.vae = AutoencoderKL.from_pretrained(snapshot, subfolder="vae")
+
+        self.unet = UNet3DConditionModel.from_pretrained_2d(
+            snapshot,
+            subfolder="unet",
+            unet_additional_kwargs={
+                "unet_use_cross_frame_attention": False,
+                "unet_use_temporal_attention": False,
+                "use_motion_module": True,
+                "motion_module_resolutions": [1, 2, 4, 8],
+                "motion_module_mid_block": False,
+                "motion_module_decoder_only": False,
+                "motion_module_type": "Vanilla",
+                "motion_module_kwargs": {
+                    "num_attention_heads": 8,
+                    "num_transformer_block": 1,
+                    "attention_block_types": ["Temporal_Self", "Temporal_Self"],
+                    "temporal_position_encoding": True,
+                    "temporal_position_encoding_max_len": 24,
+                    "temporal_attention_dim_div": 1,
+                },
+            },
+        )
+
+        assert is_xformers_available()
+        self.unet.enable_xformers_memory_efficient_attention()

minimal_animatediff/utils.py

@@ -0,0 +1,25 @@
+import os
+
+from huggingface_hub import snapshot_download
+
+
+def populate_snapshot(path: str) -> str:
+    """
+    Populates the snapshot of a specified path from the 'threeal/AnimateDiffMirrors' repository.
+
+    Args:
+        path (str): The path to the file or directory in the repository to be populated.
+
+    Returns:
+        str: The local path of the populated snapshot.
+
+    Example:
+        >>> populate_snapshot("MotionModules/mm_sd_v15.ckpt")
+        'snapshots/MotionModules/mm_sd_v15.ckpt'
+    """
+    snapshot_download(
+        repo_id="threeal/AnimateDiffMirrors",
+        local_dir="snapshots",
+        allow_patterns=[path],
+    )
+    return os.path.join("snapshots", path)

scripts/app.py

@@ -0,0 +1,48 @@
+import gradio
+import torch
+
+from deps.AnimateDiff.animatediff.utils.util import save_videos_grid
+from minimal_animatediff.animation_pipeline import create_animation_pipeline
+
+
+if __name__ == "__main__":
+    pipeline = create_animation_pipeline()
+
+    def generate(prompt, neg_prompt, steps, width, height, frames, progress=gradio.Progress()):
+        torch.manual_seed(16372571278361863751)
+
+        def update(i, _, __):
+            progress(i / float(steps + 1), desc="Sampling")
+
+        progress(0, desc="Sampling")
+        sample = pipeline(
+            prompt,
+            negative_prompt=neg_prompt,
+            num_inference_steps=int(steps),
+            guidance_scale=7.5,
+            width=int(width),
+            height=int(height),
+            video_length=int(frames),
+            callback=update,
+        ).videos
+
+        progress(steps / float(steps + 1), desc="Converting")
+        save_videos_grid(sample, "samples/sample.mp4")
+
+        progress(1)
+        return "samples/sample.mp4"
+
+    app = gradio.Interface(
+        fn=generate,
+        inputs=[
+            gradio.Textbox(label="Prompt"),
+            gradio.Textbox(label="Negative Prompt"),
+            gradio.Number(label="Steps", value=25, minimum=1),
+            gradio.Number(label="Width", value=512, minimum=1),
+            gradio.Number(label="Height", value=512, minimum=1),
+            gradio.Number(label="Frames", value=16, minimum=1),
+        ],
+        outputs=gradio.Video(label="Generated Animation"),
+    )
+
+    app.launch(share=True, enable_queue=True)

tests/test_animate.py

@@ -0,0 +1,52 @@
+from filehash import FileHash
+import pytest
+import torch
+import os
+
+from deps.AnimateDiff.animatediff.utils.util import save_videos_grid
+from minimal_animatediff.animation_pipeline import create_animation_pipeline
+
+hasher = FileHash("md5")
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+PIPELINE = None
+SAMPLE = None
+
+
+@pytest.mark.dependency()
+def test_create_animation_pipeline():
+    global PIPELINE
+    PIPELINE = create_animation_pipeline()
+
+
+@pytest.mark.dependency(depends=["test_create_animation_pipeline"])
+def test_run_animation_pipeline():
+    global SAMPLE
+    assert PIPELINE is not None
+
+    torch.manual_seed(16372571278361863751)
+    SAMPLE = PIPELINE(
+        "best quality, masterpiece, 1girl, cloudy sky, dandelion, alternate hairstyle,",
+        negative_prompt="",
+        num_inference_steps=10,
+        guidance_scale=7.5,
+        width=512,
+        height=512,
+        video_length=16,
+    ).videos
+
+
+@pytest.mark.dependency(depends=["test_run_animation_pipeline"])
+def test_save_animation_gif():
+    assert SAMPLE is not None
+    gif_path = os.path.join(dir_path, "samples/sample.gif")
+    save_videos_grid(SAMPLE, gif_path)
+    assert hasher.hash_file(gif_path) == "30cc5fb2a6446f0849889b7a84ec1c42"
+
+
+@pytest.mark.dependency(depends=["test_run_animation_pipeline"])
+def test_save_animation_mp4():
+    assert SAMPLE is not None
+    mp4_path = os.path.join(dir_path, "samples/sample.mp4")
+    save_videos_grid(SAMPLE, mp4_path)
+    assert hasher.hash_file(mp4_path) == "76494406baff329aa68d9c22ad0436cc"

tests/test_dream_booth.py

@@ -0,0 +1,19 @@
+from minimal_animatediff.dream_booth import DreamBoothModel
+
+
+def test_init_dream_booth_model():
+    model = DreamBoothModel("toonyou_beta3.safetensors")
+    assert len(model.states.keys()) == 1133
+
+
+def test_init_other_dream_booth_model():
+    model = DreamBoothModel("lyriel_v16.safetensors")
+    assert len(model.states.keys()) == 1131
+
+
+def test_init_non_existing_dream_booth_model():
+    try:
+        DreamBoothModel("invalid.safetensors")
+        assert False
+    except FileNotFoundError:
+        pass

tests/test_motion_module.py

@@ -0,0 +1,19 @@
+from minimal_animatediff.motion_module import MotionModule
+
+
+def test_init_motion_module():
+    model = MotionModule("mm_sd_v15.ckpt")
+    assert len(model.states.keys()) == 560
+
+
+def test_init_other_motion_module():
+    model = MotionModule("mm_sd_v14.ckpt")
+    assert len(model.states.keys()) == 560
+
+
+def test_init_non_existing_motion_module():
+    try:
+        MotionModule("invalid.ckpt")
+        assert False
+    except FileNotFoundError:
+        pass