Update main.py

main.py

@@ -1,8 +1,7 @@
 import torch
 import uuid
-import os
 
-from diffusers import AnimateDiffPipeline, EulerDiscreteScheduler
+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
@@ -31,9 +30,9 @@ motions = {
     "Roll left": "guoyww/animatediff-motion-lora-rolling-anticlockwise",
     "Roll right": "guoyww/animatediff-motion-lora-rolling-clockwise",
 }
-steps = [4,8]  # Different steps you want to pre-load
-models = {}
-motions_loaded = {}
+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():
@@ -41,36 +40,29 @@ if not torch.cuda.is_available():
 
 device = "cuda"
 dtype = torch.float16
-
-# 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)
+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
 
+
 @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
@@ -78,26 +70,34 @@ def generate_image(request: GenerateImageRequest):
 
     print(prompt, base, step)
 
-    if base not in models or step not in models[base]:
-        raise HTTPException(status_code=400, detail="Invalid base model or 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
 
-    pipe = models[base][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:
-        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:
+    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)
+    uvicorn.run(app, host="0.0.0.0", port=7860)