Update main.py

main.py

@@ -0,0 +1,85 @@
+import torch
+import os
+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 transformers import CLIPFeatureExtractor
+from fastapi import FastAPI, Form, HTTPException
+from fastapi.responses import FileResponse
+
+app = FastAPI()
+
+# Constants
+bases = {
+    "Cartoon": "frankjoshua/toonyou_beta6",
+    "Realistic": "emilianJR/epiCRealism",
+    "3d": "Lykon/DreamShaper",
+    "Anime": "Yntec/mistoonAnime2"
+}
+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
+feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
+
+# Function to generate image
+def generate_image(prompt, base="Realistic", motion="", step=8):
+    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(motion, 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 path
+
+# API Endpoint to generate video
+@app.post("/generate-video/")
+async def generate_video(
+    prompt: str = Form(...), 
+    base: str = Form("Realistic"), 
+    motion: str = Form(""), 
+    step: int = Form(8)
+):
+    try:
+        video_path = generate_image(prompt, base, motion, step)
+        return FileResponse(video_path, media_type="video/mp4")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Run the app
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)