Update main.py

main.py

@@ -1,7 +1,8 @@
 import torch
 import uuid
+import os
 
-from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
+from diffusers import AnimateDiffPipeline, EulerDiscreteScheduler
 from diffusers.utils import export_to_video
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
@@ -30,9 +31,9 @@ motions = {
     "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
+steps = [4,8]  # Different steps you want to pre-load
+models = {}
+motions_loaded = {}
 
 # Ensure model and scheduler are initialized in GPU-enabled function
 if not torch.cuda.is_available():
@@ -40,8 +41,22 @@ if not torch.cuda.is_available():
 
 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")
+
+# Load all base models and steps
+for base_name, base_repo in bases.items():
+    models[base_name] = {}
+    for step in steps:
+        repo = "ByteDance/AnimateDiff-Lightning"
+        ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
+        model = AnimateDiffPipeline.from_pretrained(base_repo, torch_dtype=dtype).to(device)
+        model.scheduler = EulerDiscreteScheduler.from_config(model.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
+        model.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
+        models[base_name][step] = model
+
+# Load all motion models
+for motion_name, motion_repo in motions.items():
+    motion_weights = hf_hub_download(motion_repo, "pytorch_model.bin")
+    motions_loaded[motion_name] = torch.load(motion_weights, map_location=device)
 
 # Safety checkers
 from transformers import CLIPFeatureExtractor
@@ -56,10 +71,6 @@ class GenerateImageRequest(BaseModel):
 
 @app.post("/generate-image")
 def generate_image(request: GenerateImageRequest):
-    global step_loaded
-    global base_loaded
-    global motion_loaded
-
     prompt = request.prompt
     base = request.base
     motion = request.motion
@@ -67,23 +78,18 @@ def generate_image(request: GenerateImageRequest):
 
     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 not in models or step not in models[base]:
+        raise HTTPException(status_code=400, detail="Invalid base model or 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
+    pipe = models[base][step]
 
-    if motion_loaded != motion:
+    if motion:
+        if motion not in motions_loaded:
+            raise HTTPException(status_code=400, detail="Invalid motion")
+        pipe.unet.load_state_dict(motions_loaded[motion], strict=False)
+        pipe.set_adapters(["motion"], [0.7])
+    else:
         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)