Update skyreelsinfer/skyreels_video_infer.py

skyreelsinfer/skyreels_video_infer.py

@@ -1,20 +1,22 @@
 import logging
-import os
+import os  # Keep os here
 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
+# DELAY ALL THESE IMPORTS:
+# import torch
+# from diffusers import HunyuanVideoTransformer3DModel
+# from diffusers import DiffusionPipeline
+# from PIL import Image
+# from transformers import LlamaModel
+
+# from . import TaskType
+# from .offload import Offload
+# from .offload import OffloadConfig
+# from .pipelines import SkyreelsVideoPipeline
 
-from . import TaskType
-from .offload import Offload
-from .offload import OffloadConfig
-from .pipelines import SkyreelsVideoPipeline
 
 logger = logging.getLogger("SkyReelsVideoInfer")
 logger.setLevel(logging.DEBUG)
@@ -29,11 +31,11 @@ logger.addHandler(console_handler)
 class SkyReelsVideoInfer:
     def __init__(
         self,
-        task_type: TaskType,
+        task_type, # No TaskType.
         model_id: str,
         quant_model: bool = True,
         is_offload: bool = True,
-        offload_config: OffloadConfig = None,
+        offload_config = None, # No OffloadConfig
         use_multiprocessing: bool = False,
     ):
         self.task_type = task_type
@@ -50,11 +52,19 @@ class SkyReelsVideoInfer:
         base_model_id: str = "hunyuanvideo-community/HunyuanVideo",
         quant_model: bool = True,
         device: str = "cpu",
-    ) -> SkyreelsVideoPipeline:
-        logger.info(f"load model model_id:{model_id} quan_model:{quant_model} device:{device}")
-
+    ):
+        # 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,
@@ -81,7 +91,10 @@ class SkyReelsVideoInfer:
         return pipe
 
     def _initialize_pipeline(self):
-        self.pipe: SkyreelsVideoPipeline = self._load_model(
+        #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="cpu"
         )
 
@@ -92,9 +105,11 @@ 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) # Get full output
+            output = self.pipe(image=image, **kwargs)
         else:
-            output = self.pipe(**kwargs)  # Get full output
-        return output.frames # Return frames directly
+            output = self.pipe(**kwargs)
+        return output.frames