Spaces:
Paused
Paused
| import torch | |
| 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 | |
| from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| from fastapi.responses import FileResponse | |
| import uvicorn | |
| app = FastAPI() | |
| # Constants | |
| bases = { | |
| "Cartoon": "frankjoshua/toonyou_beta6", | |
| "Realistic": "emilianJR/epiCRealism", | |
| "3d": "Lykon/DreamShaper", | |
| "Anime": "Yntec/mistoonAnime2" | |
| } | |
| motions = { | |
| "Zoom in": "guoyww/animatediff-motion-lora-zoom-in", | |
| "Zoom out": "guoyww/animatediff-motion-lora-zoom-out", | |
| "Tilt up": "guoyww/animatediff-motion-lora-tilt-up", | |
| "Tilt down": "guoyww/animatediff-motion-lora-tilt-down", | |
| "Pan left": "guoyww/animatediff-motion-lora-pan-left", | |
| "Pan right": "guoyww/animatediff-motion-lora-pan-right", | |
| "Roll left": "guoyww/animatediff-motion-lora-rolling-anticlockwise", | |
| "Roll right": "guoyww/animatediff-motion-lora-rolling-clockwise", | |
| } | |
| step_loaded = None | |
| base_loaded = "Realistic" | |
| 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") | |
| # Safety checkers | |
| from transformers import CLIPFeatureExtractor | |
| feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32") | |
| class GenerateImageRequest(BaseModel): | |
| prompt: str | |
| base: str = "Realistic" | |
| motion: str = "" | |
| step: int = 8 | |
| def generate_image(request: GenerateImageRequest): | |
| global step_loaded | |
| global base_loaded | |
| global motion_loaded | |
| prompt = request.prompt | |
| base = request.base | |
| motion = request.motion | |
| step = request.step | |
| 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 in motions: | |
| motion_repo = motions[motion] | |
| pipe.load_lora_weights(motion_repo, adapter_name="motion") | |
| pipe.set_adapters(["motion"], [0.7]) | |
| motion_loaded = motion | |
| output = pipe(prompt=prompt, guidance_scale=1.2, num_inference_steps=step) | |
| name = str(uuid.uuid4()).replace("-", "") | |
| path = f"/tmp/{name}.mp4" | |
| export_to_video(output.frames[0], path, fps=10) | |
| return FileResponse(path, media_type="video/mp4", filename=f"{name}.mp4") | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) |