Update skyreelsinfer/offload.py

skyreelsinfer/offload.py

@@ -1,333 +1,515 @@
-import spaces
-import gradio as gr
-import argparse
-import sys
+import functools
+import gc
 import os
-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 time
+from dataclasses import dataclass
 
 import torch
-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__)
+from diffusers.pipelines import DiffusionPipeline
+from torchao.dtypes.affine_quantized_tensor import AffineQuantizedTensor
 
 
-# --- Dummy Classes (Keep for standalone execution) ---
+@dataclass
 class OffloadConfig:
-    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 HunyuanVideoTransformer3DModel:
-    @staticmethod
-    def from_pretrained(*args, **kwargs):
-        return HunyuanVideoTransformer3DModel()
-
-    def to(self, device):
-        return self
-
-
-class SkyreelsVideoPipeline:
-    @staticmethod
-    def from_pretrained(*args, **kwargs):
-        return SkyreelsVideoPipeline()
-
-    def to(self, device):
-        return self
+    # 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 __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)
-
-        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)
-
-            # 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!
 
+class HfHook:
     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")
+        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
+
+    @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()
-            quantize_(transformer, float8_weight_only())
-            transformer.to("cpu")
+            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()
 
-        pipe = SkyreelsVideoPipeline.from_pretrained(
-            base_model_id, transformer=transformer, text_encoder=text_encoder, torch_dtype=torch.bfloat16
-        ).to("cpu")
-        pipe.vae.enable_tiling()
+        # 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()
-        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
+        return self
+
+    def add_module_to_blocks(self, model_name, blocks_name, submodule, prev_block_name):
 
-        if not torch.cuda.is_available():
-            raise RuntimeError("CUDA is not available. Cannot initialize model.")
+        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])
 
-        self.gpu_device = "cuda:0"
-        self.pipe = self._load_model(model_id=self.model_id, quant_model=self.quant_model)
+        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 = []
+        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
+        setattr(target_module, "_mm_id", model_name)
 
-        if self.is_offload:
-            pass
+        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}")
         else:
-            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
+            setattr(
+                target_module,
+                "forward",
+                functools.update_wrapper(functools.partial(check_empty_cuda_cache, target_module), previous_method),
             )
-            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,
-        }
 
-    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],
+    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),
         )
-    return demo
 
+        if not self.verboseLevel >= 1:
+            return
 
-if __name__ == "__main__":
-    demo = create_gradio_interface()
-    demo.queue().launch()
+        if module_id == None or module_id == "":
+            model_name = model._get_name()
+            print(f"Hooked in model '{model_name}' ({model_name})")