animate

app.py

@@ -35,42 +35,50 @@ from transformers import CLIPFeatureExtractor
 feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
 
 # Function 
-@spaces.CPU(duration=30,queue=False)
 def generate_image(prompt, base="Realistic", motion="", step=8, progress=gr.Progress()):
     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
-
-    progress((0, step))
-    def progress_callback(i, t, z):
-        progress((i+1, step))
-
-    output = pipe(prompt=prompt, guidance_scale=1.2, num_inference_steps=step, callback=progress_callback, callback_steps=1)
-
-    name = str(uuid.uuid4()).replace("-", "")
-    path = f"/tmp/{name}.mp4"
-    export_to_video(output.frames[0], path, fps=10)
-    return path
-
+    try:
+        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
+
+        progress((0, step))
+        def progress_callback(i, t, z):
+            progress((i+1, step))
+
+        output = pipe(
+            prompt=prompt, 
+            guidance_scale=1.2, 
+            num_inference_steps=step, 
+            callback=progress_callback, 
+            callback_steps=1
+        )
 
+        name = str(uuid.uuid4()).replace("-", "")
+        path = f"/tmp/{name}.mp4"
+        export_to_video(output.frames[0], path, fps=10)
+        return path
+    
+    except Exception as e:
+        print(f"Error during generation: {str(e)}")
+        return None
 # Gradio Interface
 with gr.Blocks(css="style.css") as demo:
     gr.HTML(