.

options/Video_model/Model.py

@@ -55,31 +55,28 @@ def Video(
 
     # Perform computation with appropriate dtype based on device
     if device == "cuda":
-        # Use float16 for GPU
-        with torch.autocast(device_type='cuda', dtype=torch.float16):
-            frames = pipeline(
-                image, height=height, width=width,
-                num_inference_steps=num_inference_steps,
-                min_guidance_scale=min_guidance_scale,
-                max_guidance_scale=max_guidance_scale,
-                num_frames=num_frames, fps=fps, motion_bucket_id=motion_bucket_id,
-                decode_chunk_size=8,
-                noise_aug_strength=0.02,
-                generator=generator,
-            ).frames[0]
-    else:
-        # Use float32 for CPU to avoid issues with half precision
-        with torch.autocast(device_type='cpu', dtype=torch.float32):
-            frames = pipeline(
-                image, height=height, width=width,
-                num_inference_steps=num_inference_steps,
-                min_guidance_scale=min_guidance_scale,
-                max_guidance_scale=max_guidance_scale,
-                num_frames=num_frames, fps=fps, motion_bucket_id=motion_bucket_id,
-                decode_chunk_size=8,
-                noise_aug_strength=0.02,
-                generator=generator,
-            ).frames[0]
+    # Use float16 for GPU
+    with torch.autocast(device_type='cuda', dtype=torch.float16):
+        frames = pipeline(
+            image, height=height, width=width,
+            num_inference_steps=num_inference_steps,
+            min_guidance_scale=min_guidance_scale,
+            max_guidance_scale=max_guidance_scale,
+            num_frames=num_frames, fps=fps, motion_bucket_id=motion_bucket_id,
+            generator=generator,
+        ).frames[0]
+else:
+    # Use bfloat16 for CPU as it's supported in torch.autocast
+    with torch.autocast(device_type='cpu', dtype=torch.bfloat16):
+        frames = pipeline(
+            image, height=height, width=width,
+            num_inference_steps=num_inference_steps,
+            min_guidance_scale=min_guidance_scale,
+            max_guidance_scale=max_guidance_scale,
+            num_frames=num_frames, fps=fps, motion_bucket_id=motion_bucket_id,
+            generator=generator,
+        ).frames[0]
+
 
     # Save the generated video
     save_video(frames, video_path, fps=fps, quality=5.0)