Update generate.py

generate.py

@@ -1,37 +1,54 @@
 import torch
-from diffusers import AnimateDiffPipeline, DDIMScheduler
+from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
+from diffusers.utils import export_to_gif
 
-# Load model only once (on import)
 def load_model():
-    # Example AnimateDiff model from HF Hub
-    model_id = "guoyww/animatediff-motion-adapter-v1-5"  
+    try:
+        # Load Motion Adapter
+        adapter = MotionAdapter.from_pretrained(
+            "guoyww/animatediff-motion-adapter-v1-5",
+            torch_dtype=torch.float16
+        )
 
-    pipe = AnimateDiffPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16
-    )
+        # Load AnimateDiff pipeline with Stable Diffusion 1.5
+        pipeline = AnimateDiffPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            motion_adapter=adapter,
+            torch_dtype=torch.float16
+        )
 
-    scheduler = DDIMScheduler.from_pretrained(
-        "runwayml/stable-diffusion-v1-5",
-        subfolder="scheduler"
-    )
+        # Use Euler scheduler (smoother animations)
+        pipeline.scheduler = EulerDiscreteScheduler.from_config(
+            pipeline.scheduler.config,
+            timestep_spacing="trailing"
+        )
 
-    pipe.scheduler = scheduler
-    pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        pipeline = pipeline.to(device)
 
-    return pipe
+        print("✅ Models loaded successfully!")
+        return pipeline
 
+    except Exception as e:
+        print(f"❌ Error during model loading: {e}")
+        raise
 
-# Global model instance (loaded once)
+# Load once globally
 pipe = load_model()
 
 
-# Generation function
-def generate(prompt: str, num_inference_steps: int = 50, guidance_scale: float = 7.5):
-    with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
-        result = pipe(
-            prompt,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale
-        )
-    return result.frames[0]  # returning first frame (you can adapt this to video/gif)
+def generate(prompt: str, num_frames: int = 16, steps: int = 25, guidance: float = 7.5, seed: int = 42, out_path: str = "output.gif"):
+    """
+    Generate an animated GIF from a text prompt.
+    """
+    generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
+    result = pipe(
+        prompt=prompt,
+        num_frames=num_frames,
+        num_inference_steps=steps,
+        guidance_scale=guidance,
+        generator=generator
+    )
+    frames = result.frames[0]
+    export_to_gif(frames, out_path)
+    return out_path