Spaces:
Running
Running
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 |
|