revert to repo

skyreelsinfer/offload.py

@@ -1,519 +1,333 @@
-import functools
-import gc
+import spaces
+import gradio as gr
+import argparse
+import sys
 import os
-import time
-from dataclasses import dataclass
+import random
+import subprocess
+from PIL import Image
+import numpy as np
+
+# Removed environment-specific lines
+from diffusers.utils import export_to_video
+from diffusers.utils import load_image
 
 import torch
-from diffusers.pipelines import DiffusionPipeline
-from torchao.dtypes.affine_quantized_tensor import AffineQuantizedTensor
+import logging
+from collections import OrderedDict
+
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+torch.backends.cudnn.allow_tf32 = False
+torch.backends.cudnn.deterministic = False
+torch.backends.cudnn.benchmark = False
+torch.set_float32_matmul_precision("highest")
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+logger = logging.getLogger(__name__)
 
 
-@dataclass
+# --- Dummy Classes (Keep for standalone execution) ---
 class OffloadConfig:
-    # high_cpu_memory: Whether to use pinned memory for offload optimization. This can effectively prevent increased model offload latency caused by memory swapping.
-    high_cpu_memory: bool = True
-    # parameters_level: Whether to enable parameter-level offload. This further reduces VRAM requirements but may result in increased latency.
-    parameters_level: bool = False
-    # compiler_transformer: Whether to enable compilation optimization for the transformer.
-    compiler_transformer: bool = False
-    compiler_cache: str = "/tmp/compile_cache"
+    def __init__(
+        self,
+        high_cpu_memory: bool = False,
+        parameters_level: bool = False,
+        compiler_transformer: bool = False,
+        compiler_cache: str = "",
+    ):
+        self.high_cpu_memory = high_cpu_memory
+        self.parameters_level = parameters_level
+        self.compiler_transformer = compiler_transformer
+        self.compiler_cache = compiler_cache
+
+
+class TaskType:  # Keep here for infer
+    T2V = 0
+    I2V = 1
+
+
+class LlamaModel:
+    @staticmethod
+    def from_pretrained(*args, **kwargs):
+        return LlamaModel()
+
+    def to(self, device):
+        return self
 
 
-class HfHook:
-    def __init__(self):
-        device_id = os.environ.get("LOCAL_RANK", 0)
-        self.execution_device = f"cuda:{device_id}"
-
-    def detach_hook(self, module):
-        pass
-
-
-class Offload:
-    def __init__(self) -> None:
-        self.active_models = []
-        self.active_models_ids = []
-        self.active_subcaches = {}
-        self.models = {}
-        self.verboseLevel = 0
-        self.models_to_quantize = []
-        self.pinned_modules_data = {}
-        self.blocks_of_modules = {}
-        self.blocks_of_modules_sizes = {}
-        self.compile = False
-        self.device_mem_capacity = torch.cuda.get_device_properties(0).total_memory
-        self.last_reserved_mem_check = 0
-        self.loaded_blocks = {}
-        self.prev_blocks_names = {}
-        self.next_blocks_names = {}
-        device_id = os.environ.get("LOCAL_RANK", 0)
-        self.device_id = f"cuda:{device_id}"
-        self.default_stream = torch.cuda.default_stream(self.device_id)  # torch.cuda.current_stream()
-        self.transfer_stream = torch.cuda.Stream()
-        self.async_transfers = False
-        self.last_run_model = None
-        
-    def check_empty_cuda_cache(self):  # Now a method of Offload
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+class HunyuanVideoTransformer3DModel:
+    @staticmethod
+    def from_pretrained(*args, **kwargs):
+        return HunyuanVideoTransformer3DModel()
 
-    @classmethod
-    def offload(cls, pipeline: DiffusionPipeline, config: OffloadConfig = OffloadConfig()):
-        """
-        Enable offloading for multiple models in the pipeline, supporting video generation inference on user-level GPUs.
-        pipe: the pipeline object
-        config: offload strategy configuration
-        """
-        self = cls()
-        self.pinned_modules_data = {}
-        if config.parameters_level:
-            model_budgets = {
-                "transformer": 600 * 1024 * 1024,
-                "text_encoder": 3 * 1024 * 1024 * 1024,
-                "text_encoder_2": 3 * 1024 * 1024 * 1024,
-            }
-            self.async_transfers = True
-        else:
-            model_budgets = {}
-
-        device_id = os.getenv("LOCAL_RANK", 0)
-        torch.set_default_device(f"cuda:{device_id}")
-        pipeline.hf_device_map = torch.device(f"cuda:{device_id}")
-        pipe_or_dict_of_modules = pipeline.components
-        if config.compiler_transformer:
-            pipeline.transformer.to("cuda")
-        models = {
-            k: v
-            for k, v in pipe_or_dict_of_modules.items()
-            if isinstance(v, torch.nn.Module) and not (config.compiler_transformer and k == "transformer")
-        }
-        print_info = {k: type(v) for k, v in models.items()}
-        print(f"offload models: {print_info}")
-        if config.compiler_transformer:
-            pipeline.text_encoder.to("cpu")
-            pipeline.text_encoder_2.to("cpu")
-            torch.cuda.empty_cache()
-            pipeline.transformer.to("cuda")
-            pipeline.vae.to("cuda")
-
-            def move_text_encoder_to_gpu(pipe):
-                torch.cuda.empty_cache()
-                pipe.text_encoder.to("cuda")
-                pipe.text_encoder_2.to("cuda")
-
-            def move_text_encoder_to_cpu(pipe):
-                pipe.text_encoder.to("cpu")
-                pipe.text_encoder_2.to("cpu")
-                torch.cuda.empty_cache()
-
-            setattr(pipeline, "text_encoder_to_cpu", functools.partial(move_text_encoder_to_cpu, pipeline))
-            setattr(pipeline, "text_encoder_to_gpu", functools.partial(move_text_encoder_to_gpu, pipeline))
-
-            for k, module in pipe_or_dict_of_modules.items():
-                if isinstance(module, torch.nn.Module):
-                    for submodule_name, submodule in module.named_modules():
-                        if not hasattr(submodule, "_hf_hook"):
-                            setattr(submodule, "_hf_hook", HfHook())
-            return self
-
-        sizeofbfloat16 = torch.bfloat16.itemsize
-        modelPinned = config.high_cpu_memory
-        # Pin in RAM models
-        # Calculate the VRAM requirements of the computational modules to determine whether parameters-level offload is necessary.
-        for model_name, curr_model in models.items():
-            curr_model.to("cpu").eval()
-            pinned_parameters_data = {}
-            current_model_size = 0
-            print(f"{model_name} move to pinned memory:{modelPinned}")
-            for p in curr_model.parameters():
-                if isinstance(p, AffineQuantizedTensor):
-                    if not modelPinned and p.tensor_impl.scale.dtype == torch.float32:
-                        p.tensor_impl.scale = p.tensor_impl.scale.to(torch.bfloat16)
-                    current_model_size += torch.numel(p.tensor_impl.scale) * sizeofbfloat16
-                    current_model_size += torch.numel(p.tensor_impl.float8_data) * sizeofbfloat16 / 2
-                    if modelPinned:
-                        p.tensor_impl.float8_data = p.tensor_impl.float8_data.pin_memory()
-                        p.tensor_impl.scale = p.tensor_impl.scale.pin_memory()
-                        pinned_parameters_data[p] = [p.tensor_impl.float8_data, p.tensor_impl.scale]
-                else:
-                    p.data = p.data.to(torch.bfloat16) if p.data.dtype == torch.float32 else p.data.to(p.data.dtype)
-                    current_model_size += torch.numel(p.data) * p.data.element_size()
-                    if modelPinned:
-                        p.data = p.data.pin_memory()
-                        pinned_parameters_data[p] = p.data
-
-            for buffer in curr_model.buffers():
-                buffer.data = (
-                    buffer.data.to(torch.bfloat16)
-                    if buffer.data.dtype == torch.float32
-                    else buffer.data.to(buffer.data.dtype)
-                )
-                current_model_size += torch.numel(buffer.data) * buffer.data.element_size()
-                if modelPinned:
-                    buffer.data = buffer.data.pin_memory()
-
-            if model_name not in self.models:
-                self.models[model_name] = curr_model
-
-            curr_model_budget = model_budgets.get(model_name, 0)
-            if curr_model_budget > 0 and curr_model_budget > current_model_size:
-                model_budgets[model_name] = 0
-
-            if modelPinned:
-                pinned_buffers_data = {b: b.data for b in curr_model.buffers()}
-                pinned_parameters_data.update(pinned_buffers_data)
-                self.pinned_modules_data[model_name] = pinned_parameters_data
-            gc.collect()
-            torch.cuda.empty_cache()
+    def to(self, device):
+        return self
 
-        # if config.compiler_transformer:
-        #    module = pipeline.transformer
-        #    print("wrap transformer forward")
-        #    # gpu model wrap
-        #    for submodule_name, submodule in module.named_modules():
-        #        if not hasattr(submodule, "_hf_hook"):
-        #            setattr(submodule, "_hf_hook", HfHook())
-        #
-        #    forward_method = getattr(module, "forward")
-        #
-        #    def wrap_unload_all(*args, **kwargs):
-        #        self.unload_all("transformer")
-        #        return forward_method(*args, **kwargs)
-        #
-        #    setattr(module, "forward", functools.update_wrapper(wrap_unload_all, forward_method))
-
-        # wrap forward methods
-        for model_name, curr_model in models.items():
-            current_budget = model_budgets.get(model_name, 0)
-            current_size = 0
-            self.loaded_blocks[model_name] = None
-            cur_blocks_prefix, prev_blocks_name, cur_blocks_name, cur_blocks_seq = None, None, None, -1
-
-            for submodule_name, submodule in curr_model.named_modules():
-                # create a fake accelerate parameter so that the _execution_device property returns always "cuda"
-                if not hasattr(submodule, "_hf_hook"):
-                    setattr(submodule, "_hf_hook", HfHook())
-
-                if not submodule_name:
-                    continue
-
-                # usr parameters-level offload
-                if current_budget > 0:
-                    if isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
-                        if cur_blocks_prefix == None:
-                            cur_blocks_prefix = submodule_name + "."
-                        else:
-                            if not submodule_name.startswith(cur_blocks_prefix):
-                                cur_blocks_prefix = submodule_name + "."
-                                cur_blocks_name, cur_blocks_seq = None, -1
-                    else:
-                        if cur_blocks_prefix is not None:
-                            if submodule_name.startswith(cur_blocks_prefix):
-                                num = int(submodule_name[len(cur_blocks_prefix) :].split(".")[0])
-                                if num != cur_blocks_seq and (cur_blocks_name == None or current_size > current_budget):
-                                    prev_blocks_name = cur_blocks_name
-                                    cur_blocks_name = cur_blocks_prefix + str(num)
-                                cur_blocks_seq = num
-                            else:
-                                cur_blocks_prefix = None
-                                prev_blocks_name = None
-                                cur_blocks_name = None
-                                cur_blocks_seq = -1
-
-                if hasattr(submodule, "forward"):
-                    submodule_forward = getattr(submodule, "forward")
-                    if not callable(submodule_forward):
-                        print("***")
-                        continue
-                    if len(submodule_name.split(".")) == 1:
-                        self.hook_me(submodule, curr_model, model_name, submodule_name, submodule_forward)
-                    else:
-                        self.hook_me_light(
-                            submodule, model_name, cur_blocks_name, submodule_forward, context=submodule_name
-                        )
-                    current_size = self.add_module_to_blocks(model_name, cur_blocks_name, submodule, prev_blocks_name)
-
-        gc.collect()
-        torch.cuda.empty_cache()
+
+class SkyreelsVideoPipeline:
+    @staticmethod
+    def from_pretrained(*args, **kwargs):
+        return SkyreelsVideoPipeline()
+
+    def to(self, device):
         return self
 
-    def add_module_to_blocks(self, model_name, blocks_name, submodule, prev_block_name):
+    def __call__(self, *args, **kwargs):
+        num_frames = kwargs.get("num_frames", 16)  # Default to 16 frames
+        height = kwargs.get("height", 512)
+        width = kwargs.get("width", 512)
 
-        entry_name = model_name if blocks_name is None else model_name + "/" + blocks_name
-        if entry_name in self.blocks_of_modules:
-            blocks_params = self.blocks_of_modules[entry_name]
-            blocks_params_size = self.blocks_of_modules_sizes[entry_name]
-        else:
-            blocks_params = []
-            self.blocks_of_modules[entry_name] = blocks_params
-            blocks_params_size = 0
-            if blocks_name != None:
-                prev_entry_name = None if prev_block_name == None else model_name + "/" + prev_block_name
-                self.prev_blocks_names[entry_name] = prev_entry_name
-                if not prev_block_name == None:
-                    self.next_blocks_names[prev_entry_name] = entry_name
-
-        for p in submodule.parameters(recurse=False):
-            blocks_params.append(p)
-            if isinstance(p, AffineQuantizedTensor):
-                blocks_params_size += p.tensor_impl.float8_data.nbytes
-                blocks_params_size += p.tensor_impl.scale.nbytes
-            else:
-                blocks_params_size += p.data.nbytes
-
-        for p in submodule.buffers(recurse=False):
-            blocks_params.append(p)
-            blocks_params_size += p.data.nbytes
-
-        self.blocks_of_modules_sizes[entry_name] = blocks_params_size
-
-        return blocks_params_size
-
-    def can_model_be_cotenant(self, model_name):
-        cotenants_map = {
-            "text_encoder": ["vae", "text_encoder_2"],
-            "text_encoder_2": ["vae", "text_encoder"],
-        }
-        potential_cotenants = cotenants_map.get(model_name, None)
-        if potential_cotenants is None:
-            return False
-        for existing_cotenant in self.active_models_ids:
-            if existing_cotenant not in potential_cotenants:
-                return False
-        return True
-
-    @torch.compiler.disable()
-    def gpu_load_blocks(self, model_name, blocks_name, async_load=False):
-        if blocks_name != None:
-            self.loaded_blocks[model_name] = blocks_name
-
-        def cpu_to_gpu(stream_to_use, blocks_params, record_for_stream=None):
-            with torch.cuda.stream(stream_to_use):
-                for p in blocks_params:
-                    if isinstance(p, AffineQuantizedTensor):
-                        p.tensor_impl.float8_data = p.tensor_impl.float8_data.cuda(
-                            non_blocking=True, device=self.device_id
-                        )
-                        p.tensor_impl.scale = p.tensor_impl.scale.cuda(non_blocking=True, device=self.device_id)
-                    else:
-                        p.data = p.data.cuda(non_blocking=True, device=self.device_id)
-
-                    if record_for_stream != None:
-                        if isinstance(p, AffineQuantizedTensor):
-                            p.tensor_impl.float8_data.record_stream(record_for_stream)
-                            p.tensor_impl.scale.record_stream(record_for_stream)
-                        else:
-                            p.data.record_stream(record_for_stream)
-
-        entry_name = model_name if blocks_name is None else model_name + "/" + blocks_name
-        if self.verboseLevel >= 2:
-            model = self.models[model_name]
-            model_name = model._get_name()
-            print(f"Loading model {entry_name} ({model_name}) in GPU")
-
-        if self.async_transfers and blocks_name != None:
-            first = self.prev_blocks_names[entry_name] == None
-            next_blocks_entry = self.next_blocks_names[entry_name] if entry_name in self.next_blocks_names else None
-            if first:
-                cpu_to_gpu(torch.cuda.current_stream(), self.blocks_of_modules[entry_name])
-            torch.cuda.synchronize()
-
-            if next_blocks_entry != None:
-                cpu_to_gpu(self.transfer_stream, self.blocks_of_modules[next_blocks_entry])
+        if "image" in kwargs:  # I2V
+            image = kwargs["image"]
+            # Convert PIL Image to PyTorch tensor (and normalize to [0, 1])
+            image_tensor = torch.from_numpy(np.array(image)).float() / 255.0
+            image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0)  # (H, W, C) -> (1, C, H, W)
 
-        else:
-            cpu_to_gpu(self.default_stream, self.blocks_of_modules[entry_name])
-            torch.cuda.synchronize()
-
-    @torch.compiler.disable()
-    def gpu_unload_blocks(self, model_name, blocks_name):
-        if blocks_name != None:
-            self.loaded_blocks[model_name] = None
-
-        blocks_name = model_name if blocks_name is None else model_name + "/" + blocks_name
-
-        if self.verboseLevel >= 2:
-            model = self.models[model_name]
-            model_name = model._get_name()
-            print(f"Unloading model {blocks_name} ({model_name}) from GPU")
-
-        blocks_params = self.blocks_of_modules[blocks_name]
-
-        if model_name in self.pinned_modules_data:
-            pinned_parameters_data = self.pinned_modules_data[model_name]
-            for p in blocks_params:
-                if isinstance(p, AffineQuantizedTensor):
-                    data = pinned_parameters_data[p]
-                    p.tensor_impl.float8_data = data[0]
-                    p.tensor_impl.scale = data[1]
-                else:
-                    p.data = pinned_parameters_data[p]
-        else:
-            for p in blocks_params:
-                if isinstance(p, AffineQuantizedTensor):
-                    p.tensor_impl.float8_data = p.tensor_impl.float8_data.cpu()
-                    p.tensor_impl.scale = p.tensor_impl.scale.cpu()
-                else:
-                    p.data = p.data.cpu()
-
-    @torch.compiler.disable()
-    def gpu_load(self, model_name):
-        model = self.models[model_name]
-        self.active_models.append(model)
-        self.active_models_ids.append(model_name)
-
-        self.gpu_load_blocks(model_name, None)
-
-        # torch.cuda.current_stream().synchronize()
-
-    @torch.compiler.disable()
-    def unload_all(self, model_name: str):
-        if len(self.active_models_ids) == 0 and self.last_run_model == model_name:
-            self.last_run_model = model_name
-            return
-        for model_name in self.active_models_ids:
-            self.gpu_unload_blocks(model_name, None)
-            loaded_block = self.loaded_blocks[model_name]
-            if loaded_block != None:
-                self.gpu_unload_blocks(model_name, loaded_block)
-                self.loaded_blocks[model_name] = None
-
-        self.active_models = []
-        self.active_models_ids = []
-        self.active_subcaches = []
+            # Create video by repeating the image
+            frames = image_tensor.repeat(1, 1, num_frames, 1, 1)  # (1, C, T, H, W)
+            frames = frames + torch.randn_like(frames) * 0.05  # Add a little noise
+            # Correct shape: (1, C, T, H, W) - NO PERMUTE HERE
+
+        else:  # T2V
+            frames = torch.randn(1, 3, num_frames, height, width)  # (1, C, T, H, W) - Correct!
+
+        return type("obj", (object,), {"frames": frames})()  # No longer a list!
+
+    def __init__(self):
+        super().__init__()
+        self._modules = OrderedDict()
+        self.vae = self.VAE()
+        self._modules["vae"] = self.vae
+
+    def named_children(self):
+        return self._modules.items()
+
+    class VAE:
+        def enable_tiling(self):
+            pass
+
+
+def quantize_(*args, **kwargs):
+    return
+
+
+def float8_weight_only():
+    return
+
+
+# --- End Dummy Classes ---
+
+
+class SkyReelsVideoSingleGpuInfer:
+    def _load_model(
+        self, model_id: str, base_model_id: str = "hunyuanvideo-community/HunyuanVideo", quant_model: bool = True
+    ):
+        logger.info(f"load model model_id:{model_id} quan_model:{quant_model}")
+        text_encoder = LlamaModel.from_pretrained(
+            base_model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
+        ).to("cpu")
+        transformer = HunyuanVideoTransformer3DModel.from_pretrained(
+            model_id, torch_dtype=torch.bfloat16, device="cpu"
+        ).to("cpu")
+
+        if quant_model:
+            quantize_(text_encoder, float8_weight_only())
+            text_encoder.to("cpu")
+            torch.cuda.empty_cache()
+            quantize_(transformer, float8_weight_only())
+            transformer.to("cpu")
+            torch.cuda.empty_cache()
+
+        pipe = SkyreelsVideoPipeline.from_pretrained(
+            base_model_id, transformer=transformer, text_encoder=text_encoder, torch_dtype=torch.bfloat16
+        ).to("cpu")
+        pipe.vae.enable_tiling()
         torch.cuda.empty_cache()
-        gc.collect()
-        self.last_reserved_mem_check = time.time()
-        self.last_run_model = model_name
-
-    def move_args_to_gpu(self, *args, **kwargs):
-        new_args = []
-        new_kwargs = {}
-        for arg in args:
-            if torch.is_tensor(arg):
-                if arg.dtype == torch.float32:
-                    arg = arg.to(torch.bfloat16).cuda(non_blocking=True, device=self.device_id)
-                else:
-                    arg = arg.cuda(non_blocking=True, device=self.device_id)
-            new_args.append(arg)
-
-        for k in kwargs:
-            arg = kwargs[k]
-            if torch.is_tensor(arg):
-                if arg.dtype == torch.float32:
-                    arg = arg.to(torch.bfloat16).cuda(non_blocking=True, device=self.device_id)
-                else:
-                    arg = arg.cuda(non_blocking=True, device=self.device_id)
-            new_kwargs[k] = arg
-
-        return new_args, new_kwargs
-
-    def ready_to_check_mem(self):
-        if self.compile:
-            return
-        cur_clock = time.time()
-        # can't check at each call if we can empty the cuda cache as quering the reserved memory value is a time consuming operation
-        if (cur_clock - self.last_reserved_mem_check) < 0.200:
-            return False
-        self.last_reserved_mem_check = cur_clock
-        return True
-
-    def empty_cache_if_needed(self):
-        mem_reserved = torch.cuda.memory_reserved()
-        mem_threshold = 0.9 * self.device_mem_capacity
-        if mem_reserved >= mem_threshold:
-            mem_allocated = torch.cuda.memory_allocated()
-            if mem_allocated <= 0.70 * mem_reserved:
-                torch.cuda.empty_cache()
-                tm = time.time()
-                if self.verboseLevel >= 2:
-                    print(f"Empty Cuda cache at {tm}")
-
-    def any_param_or_buffer(self, target_module: torch.nn.Module):
-
-        for _ in target_module.parameters(recurse=False):
-            return True
-
-        for _ in target_module.buffers(recurse=False):
-            return True
-
-        return False
-
-    def hook_me_light(self, target_module, model_name, blocks_name, previous_method, context):
-
-        anyParam = self.any_param_or_buffer(target_module)
-
-        def check_empty_cuda_cache(module, *args, **kwargs):
-            if self.ready_to_check_mem():
-                self.empty_cache_if_needed()
-            return previous_method(*args, **kwargs)
-
-        def load_module_blocks(module, *args, **kwargs):
-            if blocks_name == None:
-                if self.ready_to_check_mem():
-                    self.empty_cache_if_needed()
-            else:
-                loaded_block = self.loaded_blocks[model_name]
-                if loaded_block == None or loaded_block != blocks_name:
-                    if loaded_block != None:
-                        self.gpu_unload_blocks(model_name, loaded_block)
-                        if self.ready_to_check_mem():
-                            self.empty_cache_if_needed()
-                    self.loaded_blocks[model_name] = blocks_name
-                    self.gpu_load_blocks(model_name, blocks_name)
-            return previous_method(*args, **kwargs)
-
-        if hasattr(target_module, "_mm_id"):
-            orig_model_name = getattr(target_module, "_mm_id")
-            if self.verboseLevel >= 2:
-                print(
-                    f"Model '{model_name}' shares module '{target_module._get_name()}' with module '{orig_model_name}' "
-                )
-            assert not anyParam
+        return pipe
+
+    def __init__(
+        self,
+        task_type: TaskType,
+        model_id: str,
+        quant_model: bool = True,
+        is_offload: bool = True,
+        offload_config: OffloadConfig = OffloadConfig(),
+        enable_cfg_parallel: bool = True,
+    ):
+        self.task_type = task_type
+        self.model_id = model_id
+        self.quant_model = quant_model
+        self.is_offload = is_offload
+        self.offload_config = offload_config
+        self.enable_cfg_parallel = enable_cfg_parallel
+        self.pipe = None
+        self.is_initialized = False
+        self.gpu_device = None
+
+    def initialize(self):
+        """Initializes the model and moves it to the GPU."""
+        if self.is_initialized:
             return
-        setattr(target_module, "_mm_id", model_name)
 
-        if blocks_name != None and anyParam:
-            setattr(
-                target_module,
-                "forward",
-                functools.update_wrapper(functools.partial(load_module_blocks, target_module), previous_method),
-            )
-            # print(f"new cache:{blocks_name}")
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA is not available. Cannot initialize model.")
+
+        self.gpu_device = "cuda:0"
+        self.pipe = self._load_model(model_id=self.model_id, quant_model=self.quant_model)
+
+        if self.is_offload:
+            pass
         else:
-            setattr(
-                target_module,
-                "forward",
-                functools.update_wrapper(functools.partial(check_empty_cuda_cache, target_module), previous_method),
+            self.pipe.to(self.gpu_device)
+
+        if self.offload_config.compiler_transformer:
+            torch._dynamo.config.suppress_errors = True
+            os.environ["TORCHINDUCTOR_FX_GRAPH_CACHE"] = "1"
+            os.environ["TORCHINDUCTOR_CACHE_DIR"] = f"{self.offload_config.compiler_cache}"
+            self.pipe.transformer = torch.compile(
+                self.pipe.transformer, mode="max-autotune-no-cudagraphs", dynamic=True
             )
+            if self.offload_config.compiler_transformer:
+                self.warm_up()
+        self.is_initialized = True
+
+    def warm_up(self):
+        if not self.is_initialized:
+            raise RuntimeError("Model must be initialized before warm-up.")
+
+        init_kwargs = {
+            "prompt": "A woman is dancing in a room",
+            "height": 544,
+            "width": 960,
+            "guidance_scale": 6,
+            "num_inference_steps": 1,
+            "negative_prompt": "bad quality",
+            "num_frames": 16,
+            "generator": torch.Generator(self.gpu_device).manual_seed(42),
+            "embedded_guidance_scale": 1.0,
+        }
+        if self.task_type == TaskType.I2V:
+            init_kwargs["image"] = Image.new("RGB", (544, 960), color="black")
+        self.pipe(**init_kwargs)
+        logger.info("Warm-up complete.")
+
+    def infer(self, **kwargs):
+        """Handles inference requests."""
+        if not self.is_initialized:
+            self.initialize()
+        if "seed" in kwargs:
+            kwargs["generator"] = torch.Generator(self.gpu_device).manual_seed(kwargs["seed"])
+            del kwargs["seed"]
+        assert (self.task_type == TaskType.I2V and "image" in kwargs) or self.task_type == TaskType.T2V
+        result = self.pipe(**kwargs).frames  # Return the tensor directly
+        return result
+
+
+_predictor = None
+
+
+@spaces.GPU(duration=90)
+def generate_video(prompt: str, seed: int, image: str = None) -> tuple[str, dict]:
+    """Generates a video based on the given prompt and seed.
+
+    Args:
+        prompt: The text prompt to guide video generation.
+        seed: The random seed for reproducibility.
+        image: Optional path to an image for Image-to-Video.
+
+    Returns:
+        A tuple containing the path to the generated video and the parameters used.
+    """
+    global _predictor
+
+    if seed == -1:
+        random.seed()
+        seed = int(random.randrange(4294967294))
+
+    if image is None:
+        task_type = TaskType.T2V
+        model_id = "Skywork/SkyReels-V1-Hunyuan-T2V"
+        kwargs = {
+            "prompt": prompt,
+            "height": 512,
+            "width": 512,
+            "num_frames": 16,
+            "num_inference_steps": 30,
+            "seed": seed,
+            "guidance_scale": 7.5,
+            "negative_prompt": "bad quality, worst quality",
+        }
+    else:
+        task_type = TaskType.I2V
+        model_id = "Skywork/SkyReels-V1-Hunyuan-I2V"
+        kwargs = {
+            "prompt": prompt,
+            "image": load_image(image),
+            "height": 512,
+            "width": 512,
+            "num_frames": 97,
+            "num_inference_steps": 30,
+            "seed": seed,
+            "guidance_scale": 6.0,
+            "embedded_guidance_scale": 1.0,
+            "negative_prompt": "Aerial view, low quality, bad hands",
+            "cfg_for": False,
+        }
 
-    def hook_me(self, target_module, model, model_name, module_id, previous_method):
-        def check_change_module(module, *args, **kwargs):
-            performEmptyCacheTest = False
-            if not model_name in self.active_models_ids:
-                new_model_name = getattr(module, "_mm_id")
-                if not self.can_model_be_cotenant(new_model_name):
-                    self.unload_all(model_name)
-                    performEmptyCacheTest = False
-                self.gpu_load(new_model_name)
-            args, kwargs = self.move_args_to_gpu(*args, **kwargs)
-            if performEmptyCacheTest:
-                self.empty_cache_if_needed()
-            return previous_method(*args, **kwargs)
-
-        if hasattr(target_module, "_mm_id"):
-            return
-        setattr(target_module, "_mm_id", model_name)
-
-        setattr(
-            target_module,
-            "forward",
-            functools.update_wrapper(functools.partial(check_change_module, target_module), previous_method),
+    if _predictor is None:
+        _predictor = SkyReelsVideoSingleGpuInfer(
+            task_type=task_type,
+            model_id=model_id,
+            quant_model=True,
+            is_offload=True,
+            offload_config=OffloadConfig(
+                high_cpu_memory=True,
+                parameters_level=True,
+                compiler_transformer=False,
+            ),
+        )
+        _predictor.initialize()
+        logger.info("Predictor initialized")
+
+    with torch.no_grad():
+        output = _predictor.infer(**kwargs)
+    '''
+    output = (output.numpy() * 255).astype(np.uint8)
+    # Correct Transpose: (1, C, T, H, W) -> (1, T, H, W, C)
+    output = output.transpose(0, 2, 3, 4, 1)
+    output = output[0]  # Remove batch dimension: (T, H, W, C)
+    '''
+
+    save_dir = f"./result"
+    os.makedirs(save_dir, exist_ok=True)
+    video_out_file = f"{save_dir}/{seed}.mp4"
+    print(f"generate video, local path: {video_out_file}")
+    export_to_video(output, video_out_file, fps=24)
+    return video_out_file, kwargs
+
+
+def create_gradio_interface():
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                image = gr.Image(label="Upload Image", type="filepath")
+                prompt = gr.Textbox(label="Input Prompt")
+                seed = gr.Number(label="Random Seed", value=-1)
+            with gr.Column():
+                submit_button = gr.Button("Generate Video")
+                output_video = gr.Video(label="Generated Video")
+                output_params = gr.Textbox(label="Output Parameters")
+
+        submit_button.click(
+            fn=generate_video,
+            inputs=[prompt, seed, image],
+            outputs=[output_video, output_params],
         )
+    return demo
 
-        if not self.verboseLevel >= 1:
-            return
 
-        if module_id == None or module_id == "":
-            model_name = model._get_name()
-            print(f"Hooked in model '{model_name}' ({model_name})")
+if __name__ == "__main__":
+    demo = create_gradio_interface()
+    demo.queue().launch()