|
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 |
|
|
|
|
|
bases = { |
|
"ToonYou": "frankjoshua/toonyou_beta6", |
|
"epiCRealism": "emilianJR/epiCRealism" |
|
} |
|
step_loaded = None |
|
base_loaded = "ToonYou" |
|
motion_loaded = None |
|
|
|
|
|
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") |
|
|
|
|
|
@spaces.GPU(enable_queue=True) |
|
def generate_image(prompt, base, motion, 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 |
|
|
|
if motion_loaded != motion: |
|
pipe.unload_lora_weights() |
|
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 |
|
|
|
|
|
|
|
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)' |
|
) |
|
with gr.Row(): |
|
select_base = gr.Dropdown( |
|
label='Base model', |
|
choices=[ |
|
"ToonYou", |
|
"epiCRealism", |
|
], |
|
value=base_loaded, |
|
interactive=True |
|
) |
|
select_motion = gr.Dropdown( |
|
label='Motion LoRAs', |
|
choices=[ |
|
("None", None), |
|
("Zoom in", "v2_lora_ZoomIn.ckpt"), |
|
("Zoom out", "v2_lora_ZoomOut.ckpt"), |
|
], |
|
value=None, |
|
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() |