Taf2023/AnimateDiff-Lightning

1import gradio as gr2import torch3import os4import spaces5import uuid6 7from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler8from diffusers.utils import export_to_video9from huggingface_hub import hf_hub_download10from safetensors.torch import load_file11from PIL import Image12 13# Constants14bases = {15    "ToonYou": "frankjoshua/toonyou_beta6",16    "epiCRealism": "emilianJR/epiCRealism"17}18step_loaded = None19base_loaded = "ToonYou"20motion_loaded = None21 22# Ensure model and scheduler are initialized in GPU-enabled function23if not torch.cuda.is_available():24    raise NotImplementedError("No GPU detected!")25 26device = "cuda"27dtype = torch.float1628pipe = AnimateDiffPipeline.from_pretrained(bases[base_loaded], torch_dtype=dtype).to(device)29pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")30 31# Safety checkers32from safety_checker import StableDiffusionSafetyChecker33from transformers import CLIPFeatureExtractor34 35safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to(device)36feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")37 38def check_nsfw_images(images: list[Image.Image]) -> list[bool]:39    safety_checker_input = feature_extractor(images, return_tensors="pt").to(device)40    has_nsfw_concepts = safety_checker(images=[images], clip_input=safety_checker_input.pixel_values.to(device))41    return has_nsfw_concepts42 43# Function 44@spaces.GPU(enable_queue=True)45def generate_image(prompt, base, motion, step, progress=gr.Progress()):46    global step_loaded47    global base_loaded48    global motion_loaded49    print(prompt, base, step)50 51    if step_loaded != step:52        repo = "ByteDance/AnimateDiff-Lightning"53        ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"54        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)55        step_loaded = step56 57    if base_loaded != base:58        pipe.unet.load_state_dict(torch.load(hf_hub_download(bases[base], "unet/diffusion_pytorch_model.bin"), map_location=device), strict=False)59        base_loaded = base60 61    if motion_loaded != motion:62        pipe.unload_lora_weights()63        if motion != "":64            pipe.load_lora_weights(motion, adapter_name="motion")65            pipe.set_adapters(["motion"], [0.7])66        motion_loaded = motion67 68    progress((0, step))69    def progress_callback(i, t, z):70        progress((i+1, step))71 72    output = pipe(prompt=prompt, guidance_scale=1.0, num_inference_steps=step, callback=progress_callback, callback_steps=1)73 74    has_nsfw_concepts = check_nsfw_images([output.frames[0][0]])75    if has_nsfw_concepts[0]:76        gr.Warning("NSFW content detected.")77        return None78 79    name = str(uuid.uuid4()).replace("-", "")80    path = f"/tmp/{name}.mp4"81    export_to_video(output.frames[0], path, fps=10)82    return path83 84 85# Gradio Interface86with gr.Blocks(css="style.css") as demo:87    gr.HTML(88        "<h1><center>AnimateDiff-Lightning ⚡</center></h1>" +89        "<p><center>Lightning-fast text-to-video generation</center></p>" +90        "<p><center><a href='https://huggingface.co/ByteDance/AnimateDiff-Lightning'>https://huggingface.co/ByteDance/AnimateDiff-Lightning</a></center></p>"91    )92    with gr.Group():93        with gr.Row():94            prompt = gr.Textbox(95                label='Prompt (English)'96            )97        with gr.Row():98            select_base = gr.Dropdown(99                label='Base model',100                choices=[101                    "ToonYou", 102                    "epiCRealism",103                ],104                value=base_loaded,105                interactive=True106            )107            select_motion = gr.Dropdown(108                label='Motion',109                choices=[110                    ("Default", ""),111                    ("Zoom in", "guoyww/animatediff-motion-lora-zoom-in"),112                    ("Zoom out", "guoyww/animatediff-motion-lora-zoom-out"),113                    ("Tilt up", "guoyww/animatediff-motion-lora-tilt-up"),114                    ("Tilt down", "guoyww/animatediff-motion-lora-tilt-down"),115                    ("Pan left", "guoyww/animatediff-motion-lora-pan-left"),116                    ("Pan right", "guoyww/animatediff-motion-lora-pan-right"),117                    ("Roll left", "guoyww/animatediff-motion-lora-rolling-anticlockwise"),118                    ("Roll right", "guoyww/animatediff-motion-lora-rolling-clockwise"),119                ],120                value="",121                interactive=True122            )123            select_step = gr.Dropdown(124                label='Inference steps',125                choices=[126                    ('1-Step', 1), 127                    ('2-Step', 2),128                    ('4-Step', 4),129                    ('8-Step', 8)],130                value=4,131                interactive=True132            )133            submit = gr.Button(134                scale=1,135                variant='primary'136            )137    video = gr.Video(138        label='AnimateDiff-Lightning',139        autoplay=True,140        height=512,141        width=512,142        elem_id="video_output"143    )144 145    prompt.submit(146        fn=generate_image,147        inputs=[prompt, select_base, select_motion, select_step],148        outputs=video,149    )150    submit.click(151        fn=generate_image,152        inputs=[prompt, select_base, select_motion, select_step],153        outputs=video,154    )155 156demo.queue().launch()