cpu

options/Video_model/Model.py

@@ -3,6 +3,8 @@ from diffusers import StableVideoDiffusionPipeline
 from diffusers.utils import load_image, export_to_video
 from PIL import Image
 
+device="cuda" if torch.cuda.is_available() else "cpu"
+
 pipeline = StableVideoDiffusionPipeline.from_pretrained(
     "stabilityai/stable-video-diffusion-img2vid-xt-1-1", torch_dtype=torch.float16, variant="fp16"
 )
@@ -13,6 +15,6 @@ def Video(image):
     image = image.resize((1024, 576))
 
     generator = torch.manual_seed(42)
-    frames = pipeline(image, decode_chunk_size=8, generator=generator).frames[0]
+    frames = pipeline(image, decode_chunk_size=8, generator=generator,device=device).frames[0]
     export_to_video(frames, "generated.mp4", fps=7)
     return "generated.mp4"