File size: 3,378 Bytes
90ee73b
 
 
 
 
 
 
 
 
 
 
 
 
f3a1f2e
 
 
 
 
 
90ee73b
 
f3a1f2e
90ee73b
 
f3a1f2e
 
 
 
 
90ee73b
 
f3a1f2e
5114719
 
 
90ee73b
f3a1f2e
87d5fe9
 
 
f3a1f2e
 
 
 
 
90ee73b
87d5fe9
90ee73b
 
 
 
 
 
 
 
 
 
 
 
87d5fe9
f3a1f2e
87d5fe9
 
f3a1f2e
 
 
 
 
 
 
 
 
 
 
87d5fe9
 
 
 
 
7096730
87d5fe9
 
 
 
 
 
 
 
 
7096730
13600b0
 
87d5fe9
90ee73b
 
 
f3a1f2e
90ee73b
 
 
 
f3a1f2e
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
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"

# 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, step):
    global step_loaded
    global base_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

    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>")
    gr.HTML("<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)',
                scale=8
            )
            select_base = gr.Dropdown(
                label='Base model',
                choices=[
                    "ToonYou", 
                    "epiCRealism",
                ],
                value=base_loaded,
                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_step],
        outputs=video,
    )
    submit.click(
        fn=generate_image,
        inputs=[prompt, select_base, select_step],
        outputs=video,
    )

demo.queue().launch()