Update skyreelsinfer/skyreels_video_infer.py

skyreelsinfer/skyreels_video_infer.py

@@ -9,9 +9,6 @@ import torch
 from diffusers import HunyuanVideoTransformer3DModel
 from diffusers import DiffusionPipeline
 from PIL import Image
-# DELAY torchao imports:
-# from torchao.quantization import float8_weight_only
-# from torchao.quantization import quantize_
 from transformers import LlamaModel
 
 from . import TaskType
@@ -19,7 +16,6 @@ from .offload import Offload
 from .offload import OffloadConfig
 from .pipelines import SkyreelsVideoPipeline
 
-
 logger = logging.getLogger("SkyReelsVideoInfer")
 logger.setLevel(logging.DEBUG)
 console_handler = logging.StreamHandler()
@@ -30,8 +26,6 @@ formatter = logging.Formatter(
 console_handler.setFormatter(formatter)
 logger.addHandler(console_handler)
 
-
-
 class SkyReelsVideoInfer:
     def __init__(
         self,
@@ -55,11 +49,10 @@ class SkyReelsVideoInfer:
         model_id: str,
         base_model_id: str = "hunyuanvideo-community/HunyuanVideo",
         quant_model: bool = True,
-        device: str = "cpu",  # Use string "cpu"
+        device: str = "cpu",
     ) -> SkyreelsVideoPipeline:
         logger.info(f"load model model_id:{model_id} quan_model:{quant_model} device:{device}")
 
-        # DELAYED IMPORTS:
         from torchao.quantization import float8_weight_only
         from torchao.quantization import quantize_
 
@@ -87,7 +80,6 @@ class SkyReelsVideoInfer:
         pipe.vae.enable_tiling()
         return pipe
 
-
     def _initialize_pipeline(self):
         self.pipe: SkyreelsVideoPipeline = self._load_model(
             model_id=self.model_id, quant_model=self.quant_model, device="cpu"
@@ -99,11 +91,10 @@ class SkyReelsVideoInfer:
                 config=self.offload_config,
             )
 
-
     def inference(self, kwargs):
         if self.task_type == TaskType.I2V:
             image = kwargs.pop("image")
-            output = self.pipe(image=image, **kwargs).frames
+            output = self.pipe(image=image, **kwargs) # Get full output
         else:
-            output = self.pipe(**kwargs).frames
-        return output
+            output = self.pipe(**kwargs)  # Get full output
+        return output.frames # Return frames directly