Update skyreelsinfer/skyreels_video_infer.py

skyreelsinfer/skyreels_video_infer.py

@@ -4,6 +4,17 @@ import time
 from datetime import timedelta
 from typing import Any
 from typing import Dict
+import torch
+from diffusers import HunyuanVideoTransformer3DModel
+from diffusers import DiffusionPipeline
+from PIL import Image
+from transformers import LlamaModel
+from torchao.quantization import float8_weight_only
+from torchao.quantization import quantize_
+from .pipelines import SkyreelsVideoPipeline  # Local import
+from .offload import Offload
+from .offload import OffloadConfig
+from . import TaskType
 
 # DELAY ALL THESE IMPORTS:
 # import torch
@@ -34,7 +45,7 @@ class SkyReelsVideoInfer:
         model_id: str,
         quant_model: bool = True,
         is_offload: bool = True,
-        offload_config = None, # No OffloadConfig
+        offload_config: OffloadConfig = OffloadConfig(),
         use_multiprocessing: bool = False,
     ):
         self.task_type = task_type
@@ -42,7 +53,6 @@ class SkyReelsVideoInfer:
         self.quant_model = quant_model
         self.is_offload = is_offload
         self.offload_config = offload_config
-
         self._initialize_pipeline()
 
     def _load_model(
@@ -52,31 +62,16 @@ class SkyReelsVideoInfer:
         quant_model: bool = True,
         device: str = "cuda",
     ):
-        # DELAYED IMPORTS:
-        import torch
-        from diffusers import HunyuanVideoTransformer3DModel
-        from diffusers import DiffusionPipeline
-        from PIL import Image
-        from transformers import LlamaModel
-        from torchao.quantization import float8_weight_only
-        from torchao.quantization import quantize_
-        from .pipelines import SkyreelsVideoPipeline  # Local import
-
-
         logger.info(f"load model model_id:{model_id} quan_model:{quant_model} device:{device}")
-
         text_encoder = LlamaModel.from_pretrained(
             base_model_id,
             subfolder="text_encoder",
             torch_dtype=torch.bfloat16,
         ).to(device)
-
         transformer = HunyuanVideoTransformer3DModel.from_pretrained(
             model_id,
             torch_dtype=torch.bfloat16,
         ).to(device)
-
-
         if quant_model:
             quantize_(text_encoder, float8_weight_only(), device=device)
             quantize_(transformer, float8_weight_only(), device=device)
@@ -90,13 +85,9 @@ class SkyReelsVideoInfer:
         return pipe
 
     def _initialize_pipeline(self):
-        #More Delayed Imports
-        from .offload import Offload
-
         self.pipe = self._load_model( #No : SkyreelsVideoPipeline
             model_id=self.model_id, quant_model=self.quant_model, device="cuda"
         )
-
         if self.is_offload and self.offload_config:
             Offload.offload(
                 pipeline=self.pipe,
@@ -104,8 +95,6 @@ class SkyReelsVideoInfer:
             )
 
     def inference(self, kwargs):
-        #DELAYED IMPORTS
-        from . import TaskType
         if self.task_type == TaskType.I2V:
             image = kwargs.pop("image")
             output = self.pipe(image=image, **kwargs)