Spaces:
Sleeping
Sleeping
Update skyreels_v2_infer/pipelines/image2video_pipeline.py
Browse files
skyreels_v2_infer/pipelines/image2video_pipeline.py
CHANGED
|
@@ -117,7 +117,7 @@ class Image2VideoPipeline:
|
|
| 117 |
)
|
| 118 |
|
| 119 |
self.transformer.to(self.device)
|
| 120 |
-
with torch.
|
| 121 |
self.scheduler.set_timesteps(num_inference_steps, device=self.device, shift=shift)
|
| 122 |
timesteps = self.scheduler.timesteps
|
| 123 |
|
|
|
|
| 117 |
)
|
| 118 |
|
| 119 |
self.transformer.to(self.device)
|
| 120 |
+
with torch.amp.autocast("cuda", dtype=self.transformer.dtype), torch.no_grad():
|
| 121 |
self.scheduler.set_timesteps(num_inference_steps, device=self.device, shift=shift)
|
| 122 |
timesteps = self.scheduler.timesteps
|
| 123 |
|