Update skyreelsinfer/skyreels_video_infer.py

skyreelsinfer/skyreels_video_infer.py

@@ -38,12 +38,12 @@ class SkyReelsVideoSingleGpuInfer:
         text_encoder = LlamaModel.from_pretrained(
             base_model_id,
             subfolder="text_encoder",
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
         ).to("cpu")
         transformer = HunyuanVideoTransformer3DModel.from_pretrained(
             model_id,
             # subfolder="transformer",
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
             device="cpu",
         ).to("cpu")
         if quant_model:
@@ -57,7 +57,7 @@ class SkyReelsVideoSingleGpuInfer:
             base_model_id,
             transformer=transformer,
             text_encoder=text_encoder,
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
         ).to("cpu")
         pipe.vae.enable_tiling()
         torch.cuda.empty_cache()