File size: 4,104 Bytes
90ee73b
 
 
 
 
 
 
 
 
 
 
 
 
f3a1f2e
 
 
 
 
 
337bc14
90ee73b
 
f3a1f2e
90ee73b
 
f3a1f2e
 
 
 
 
90ee73b
 
337bc14
5114719
 
1514a70
5114719
90ee73b
f3a1f2e
87d5fe9
 
 
f3a1f2e
 
 
 
 
90ee73b
337bc14
 
1514a70
 
 
337bc14
 
87d5fe9
90ee73b
 
 
 
 
 
 
 
24e347a
 
 
 
 
90ee73b
 
87d5fe9
337bc14
87d5fe9
337bc14
f3a1f2e
 
 
 
 
 
 
 
 
337bc14
27f6e5d
337bc14
27f6e5d
337bc14
 
 
1514a70
337bc14
 
f3a1f2e
 
87d5fe9
 
 
 
 
7096730
87d5fe9
 
 
 
 
 
 
 
 
7096730
13600b0
 
87d5fe9
90ee73b
 
 
337bc14
90ee73b
 
 
 
337bc14
90ee73b
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import gradio as gr
import torch
import os
import spaces
import uuid

from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
from diffusers.utils import export_to_video
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from PIL import Image

# Constants
bases = {
    "ToonYou": "frankjoshua/toonyou_beta6",
    "epiCRealism": "emilianJR/epiCRealism"
}
step_loaded = None
base_loaded = "ToonYou"
motion_loaded = None

# Ensure model and scheduler are initialized in GPU-enabled function
if not torch.cuda.is_available():
    raise NotImplementedError("No GPU detected!")

device = "cuda"
dtype = torch.float16
pipe = AnimateDiffPipeline.from_pretrained(bases[base_loaded], torch_dtype=dtype).to(device)
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")

# Function 
@spaces.GPU(enable_queue=True)
def generate_image(prompt, base, motion, step):
    global step_loaded
    global base_loaded
    global motion_loaded
    print(prompt, base, step)

    if step_loaded != step:
        repo = "ByteDance/AnimateDiff-Lightning"
        ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
        step_loaded = step

    if base_loaded != base:
        pipe.unet.load_state_dict(torch.load(hf_hub_download(bases[base], "unet/diffusion_pytorch_model.bin"), map_location=device), strict=False)
        base_loaded = base

    if motion_loaded != motion:
        pipe.unload_lora_weights()
        if motion != "":
            pipe.load_lora_weights(hf_hub_download("guoyww/animatediff", motion), adapter_name="motion")
            pipe.set_adapters(["motion"], [0.7])
        motion_loaded = motion

    output = pipe(prompt=prompt, guidance_scale=1.0, num_inference_steps=step)
    name = str(uuid.uuid4()).replace("-", "")
    path = f"/tmp/{name}.mp4"
    export_to_video(output.frames[0], path, fps=10)
    return path


# Gradio Interface
with gr.Blocks(css="style.css") as demo:
    gr.HTML(
        "<h1><center>AnimateDiff-Lightning ⚡</center></h1>" +
        "<p><center>Lightning-fast text-to-video generation</center></p>" +
        "<p><center><a href='https://huggingface.co/ByteDance/AnimateDiff-Lightning'>https://huggingface.co/ByteDance/AnimateDiff-Lightning</a></center></p>"
    )
    with gr.Group():
        with gr.Row():
            prompt = gr.Textbox(
                label='Prompt (English)'
            )
        with gr.Row():
            select_base = gr.Dropdown(
                label='Base model',
                choices=[
                    "ToonYou", 
                    "epiCRealism",
                ],
                value=base_loaded,
                interactive=True
            )
            select_motion = gr.Dropdown(
                label='Motion',
                choices=[
                    ("Default", ""),
                    ("Zoom in", "v2_lora_ZoomIn.ckpt"),
                    ("Zoom out", "v2_lora_ZoomOut.ckpt"),
                ],
                value="",
                interactive=True
            )
            select_step = gr.Dropdown(
                label='Inference steps',
                choices=[
                    ('1-Step', 1), 
                    ('2-Step', 2),
                    ('4-Step', 4),
                    ('8-Step', 8)],
                value=4,
                interactive=True
            )
            submit = gr.Button(
                scale=1,
                variant='primary'
            )
    video = gr.Video(
        label='AnimateDiff-Lightning',
        autoplay=True,
        height=512,
        width=512,
        elem_id="video_output"
    )

    prompt.submit(
        fn=generate_image,
        inputs=[prompt, select_base, select_motion, select_step],
        outputs=video,
    )
    submit.click(
        fn=generate_image,
        inputs=[prompt, select_base, select_motion, select_step],
        outputs=video,
    )

demo.queue().launch()