Upload 5 files

skyreelsinfer/__init__.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class TaskType(str, Enum):
+    T2V = "text2video"
+    I2V = "image2video"

skyreelsinfer/offload.py

@@ -0,0 +1,515 @@
+import functools
+import gc
+import os
+import time
+from dataclasses import dataclass
+
+import torch
+from diffusers.pipelines import DiffusionPipeline
+from torchao.dtypes.affine_quantized_tensor import AffineQuantizedTensor
+
+
+@dataclass
+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"
+
+
+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
+
+    @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()
+
+        # 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 self
+
+    def add_module_to_blocks(self, model_name, blocks_name, submodule, prev_block_name):
+
+        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])
+
+        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 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:
+            setattr(
+                target_module,
+                "forward",
+                functools.update_wrapper(functools.partial(check_empty_cuda_cache, target_module), previous_method),
+            )
+
+    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 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})")

skyreelsinfer/pipelines/__init__.py

@@ -0,0 +1 @@
+from .pipeline_skyreels_video import SkyreelsVideoPipeline

skyreelsinfer/pipelines/pipeline_skyreels_video.py

@@ -0,0 +1,425 @@
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+import numpy as np
+import torch
+from diffusers import HunyuanVideoPipeline
+from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import DEFAULT_PROMPT_TEMPLATE
+from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import HunyuanVideoPipelineOutput
+from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import MultiPipelineCallbacks
+from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import PipelineCallback
+from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import retrieve_timesteps
+from PIL import Image
+
+
+def resizecrop(image, th, tw):
+    w, h = image.size
+    if h / w > th / tw:
+        new_w = int(w)
+        new_h = int(new_w * th / tw)
+    else:
+        new_h = int(h)
+        new_w = int(new_h * tw / th)
+    left = (w - new_w) / 2
+    top = (h - new_h) / 2
+    right = (w + new_w) / 2
+    bottom = (h + new_h) / 2
+    image = image.crop((left, top, right, bottom))
+    return image
+
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+class SkyreelsVideoPipeline(HunyuanVideoPipeline):
+    """
+    support i2v and t2v
+    support true_cfg
+    """
+
+    @property
+    def guidance_rescale(self):
+        return self._guidance_rescale
+
+    @property
+    def clip_skip(self):
+        return self._clip_skip
+
+    @property
+    def do_classifier_free_guidance(self):
+        # return self._guidance_scale > 1 and self.transformer.config.time_cond_proj_dim is None
+        return self._guidance_scale > 1
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        do_classifier_free_guidance: bool,
+        negative_prompt: str = "",
+        prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+        num_videos_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_attention_mask: Optional[torch.Tensor] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        max_sequence_length: int = 256,
+    ):
+        num_hidden_layers_to_skip = self.clip_skip if self.clip_skip is not None else 0
+        print(f"num_hidden_layers_to_skip: {num_hidden_layers_to_skip}")
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds(
+                prompt,
+                prompt_template,
+                num_videos_per_prompt,
+                device=device,
+                dtype=dtype,
+                num_hidden_layers_to_skip=num_hidden_layers_to_skip,
+                max_sequence_length=max_sequence_length,
+            )
+        if negative_prompt_embeds is None and do_classifier_free_guidance:
+            negative_prompt_embeds, negative_attention_mask = self._get_llama_prompt_embeds(
+                negative_prompt,
+                prompt_template,
+                num_videos_per_prompt,
+                device=device,
+                dtype=dtype,
+                num_hidden_layers_to_skip=num_hidden_layers_to_skip,
+                max_sequence_length=max_sequence_length,
+            )
+        if self.text_encoder_2 is not None and pooled_prompt_embeds is None:
+            pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                prompt,
+                num_videos_per_prompt,
+                device=device,
+                dtype=dtype,
+                max_sequence_length=77,
+            )
+            if negative_pooled_prompt_embeds is None and do_classifier_free_guidance:
+                negative_pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                    negative_prompt,
+                    num_videos_per_prompt,
+                    device=device,
+                    dtype=dtype,
+                    max_sequence_length=77,
+                )
+        return (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_attention_mask,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        )
+
+    def image_latents(
+        self,
+        initial_image,
+        batch_size,
+        height,
+        width,
+        device,
+        dtype,
+        num_channels_latents,
+        video_length,
+    ):
+        initial_image = initial_image.unsqueeze(2)
+        image_latents = self.vae.encode(initial_image).latent_dist.sample()
+        if hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor:
+            image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+        else:
+            image_latents = image_latents * self.vae.config.scaling_factor
+        padding_shape = (
+            batch_size,
+            num_channels_latents,
+            video_length - 1,
+            int(height) // self.vae_scale_factor_spatial,
+            int(width) // self.vae_scale_factor_spatial,
+        )
+        latent_padding = torch.zeros(padding_shape, device=device, dtype=dtype)
+        image_latents = torch.cat([image_latents, latent_padding], dim=2)
+        return image_latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: str,
+        negative_prompt: str = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion",
+        height: int = 720,
+        width: int = 1280,
+        num_frames: int = 129,
+        num_inference_steps: int = 50,
+        sigmas: List[float] = None,
+        guidance_scale: float = 1.0,
+        num_videos_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_attention_mask: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        clip_skip: Optional[int] = 2,
+        callback_on_step_end: Optional[
+            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+        ] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
+        max_sequence_length: int = 256,
+        embedded_guidance_scale: Optional[float] = 6.0,
+        image: Optional[Union[torch.Tensor, Image.Image]] = None,
+        cfg_for: bool = False,
+    ):
+        if hasattr(self, "text_encoder_to_gpu"):
+            self.text_encoder_to_gpu()
+
+        if image is not None and isinstance(image, Image.Image):
+            image = resizecrop(image, height, width)
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            None,
+            height,
+            width,
+            prompt_embeds,
+            callback_on_step_end_tensor_inputs,
+            prompt_template,
+        )
+        #  add negative prompt check
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._attention_kwargs = attention_kwargs
+        self._interrupt = False
+
+        device = self._execution_device
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        # 3. Encode input prompt
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_attention_mask,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            prompt_template=prompt_template,
+            num_videos_per_prompt=num_videos_per_prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_attention_mask=negative_attention_mask,
+            device=device,
+            max_sequence_length=max_sequence_length,
+        )
+
+        transformer_dtype = self.transformer.dtype
+        prompt_embeds = prompt_embeds.to(transformer_dtype)
+        prompt_attention_mask = prompt_attention_mask.to(transformer_dtype)
+        if pooled_prompt_embeds is not None:
+            pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype)
+
+        ## Embeddings are concatenated to form a batch.
+        if self.do_classifier_free_guidance:
+            negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
+            negative_attention_mask = negative_attention_mask.to(transformer_dtype)
+            if negative_pooled_prompt_embeds is not None:
+                negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to(transformer_dtype)
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+            if prompt_attention_mask is not None:
+                prompt_attention_mask = torch.cat([negative_attention_mask, prompt_attention_mask])
+            if pooled_prompt_embeds is not None:
+                pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds])
+
+        # 4. Prepare timesteps
+        sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+        )
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels
+        if image is not None:
+            num_channels_latents = int(num_channels_latents / 2)
+            image = self.video_processor.preprocess(image, height=height, width=width).to(
+                device, dtype=prompt_embeds.dtype
+            )
+        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            num_latent_frames,
+            torch.float32,
+            device,
+            generator,
+            latents,
+        )
+        # add image latents
+        if image is not None:
+            image_latents = self.image_latents(
+                image, batch_size, height, width, device, torch.float32, num_channels_latents, num_latent_frames
+            )
+
+            image_latents = image_latents.to(transformer_dtype)
+        else:
+            image_latents = None
+
+        # 6. Prepare guidance condition
+        if self.do_classifier_free_guidance:
+            guidance = (
+                torch.tensor([embedded_guidance_scale] * latents.shape[0] * 2, dtype=transformer_dtype, device=device)
+                * 1000.0
+            )
+        else:
+            guidance = (
+                torch.tensor([embedded_guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device)
+                * 1000.0
+            )
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        if hasattr(self, "text_encoder_to_cpu"):
+            self.text_encoder_to_cpu()
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                latents = latents.to(transformer_dtype)
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                # timestep = t.expand(latents.shape[0]).to(latents.dtype)
+                if image_latents is not None:
+                    latent_image_input = (
+                        torch.cat([image_latents] * 2) if self.do_classifier_free_guidance else image_latents
+                    )
+                    latent_model_input = torch.cat([latent_model_input, latent_image_input], dim=1)
+                timestep = t.repeat(latent_model_input.shape[0]).to(torch.float32)
+                if cfg_for and self.do_classifier_free_guidance:
+                    noise_pred_list = []
+                    for idx in range(latent_model_input.shape[0]):
+                        noise_pred_uncond = self.transformer(
+                            hidden_states=latent_model_input[idx].unsqueeze(0),
+                            timestep=timestep[idx].unsqueeze(0),
+                            encoder_hidden_states=prompt_embeds[idx].unsqueeze(0),
+                            encoder_attention_mask=prompt_attention_mask[idx].unsqueeze(0),
+                            pooled_projections=pooled_prompt_embeds[idx].unsqueeze(0),
+                            guidance=guidance[idx].unsqueeze(0),
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
+                        )[0]
+                        noise_pred_list.append(noise_pred_uncond)
+                    noise_pred = torch.cat(noise_pred_list, dim=0)
+                else:
+                    noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep,
+                        encoder_hidden_states=prompt_embeds,
+                        encoder_attention_mask=prompt_attention_mask,
+                        pooled_projections=pooled_prompt_embeds,
+                        guidance=guidance,
+                        attention_kwargs=attention_kwargs,
+                        return_dict=False,
+                    )[0]
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(
+                        noise_pred,
+                        noise_pred_text,
+                        guidance_rescale=self.guidance_rescale,
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+        if not output_type == "latent":
+            latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor
+            video = self.vae.decode(latents, return_dict=False)[0]
+            video = self.video_processor.postprocess_video(video, output_type=output_type)
+        else:
+            video = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (video,)
+
+        return HunyuanVideoPipelineOutput(frames=video)

skyreelsinfer/skyreels_video_infer.py

@@ -0,0 +1,258 @@
+import logging
+import os
+import threading
+import time
+from datetime import timedelta
+from typing import Any
+from typing import Dict
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from diffusers import HunyuanVideoTransformer3DModel
+from PIL import Image
+from torchao.quantization import float8_weight_only
+from torchao.quantization import quantize_
+from transformers import LlamaModel
+
+from . import TaskType
+from .offload import Offload
+from .offload import OffloadConfig
+from .pipelines import SkyreelsVideoPipeline
+
+logger = logging.getLogger("SkyreelsVideoInfer")
+logger.setLevel(logging.DEBUG)
+console_handler = logging.StreamHandler()
+console_handler.setLevel(logging.DEBUG)
+formatter = logging.Formatter(
+    f"%(asctime)s - %(name)s - %(levelname)s - [%(filename)s:%(lineno)d - %(funcName)s] - %(message)s"
+)
+console_handler.setFormatter(formatter)
+logger.addHandler(console_handler)
+
+
+class SkyReelsVideoSingleGpuInfer:
+    def _load_model(
+        self,
+        model_id: str,
+        base_model_id: str = "hunyuanvideo-community/HunyuanVideo",
+        quant_model: bool = True,
+        gpu_device: str = "cuda:0",
+    ) -> SkyreelsVideoPipeline:
+        logger.info(f"load model model_id:{model_id} quan_model:{quant_model} gpu_device:{gpu_device}")
+        text_encoder = LlamaModel.from_pretrained(
+            base_model_id,
+            subfolder="text_encoder",
+            torch_dtype=torch.bfloat16,
+        ).to("cpu")
+        transformer = HunyuanVideoTransformer3DModel.from_pretrained(
+            model_id,
+            # subfolder="transformer",
+            torch_dtype=torch.bfloat16,
+            device="cpu",
+        ).to("cpu")
+        if quant_model:
+            quantize_(text_encoder, float8_weight_only(), device=gpu_device)
+            text_encoder.to("cpu")
+            torch.cuda.empty_cache()
+            quantize_(transformer, float8_weight_only(), device=gpu_device)
+            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()
+        return pipe
+
+    def __init__(
+        self,
+        task_type: TaskType,
+        model_id: str,
+        quant_model: bool = True,
+        local_rank: int = 0,
+        world_size: int = 1,
+        is_offload: bool = True,
+        offload_config: OffloadConfig = OffloadConfig(),
+        enable_cfg_parallel: bool = True,
+    ):
+        self.task_type = task_type
+        self.gpu_rank = local_rank
+        dist.init_process_group(
+            backend="nccl",
+            init_method="tcp://127.0.0.1:23456",
+            timeout=timedelta(seconds=600),
+            world_size=world_size,
+            rank=local_rank,
+        )
+        os.environ["LOCAL_RANK"] = str(local_rank)
+        logger.info(f"rank:{local_rank} Distributed backend: {dist.get_backend()}")
+        torch.cuda.set_device(dist.get_rank())
+        torch.backends.cuda.enable_cudnn_sdp(False)
+        gpu_device = f"cuda:{dist.get_rank()}"
+
+        self.pipe: SkyreelsVideoPipeline = self._load_model(
+            model_id=model_id, quant_model=quant_model, gpu_device=gpu_device
+        )
+
+        from para_attn.context_parallel import init_context_parallel_mesh
+        from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
+        from para_attn.parallel_vae.diffusers_adapters import parallelize_vae
+
+        max_batch_dim_size = 2 if enable_cfg_parallel and world_size > 1 else 1
+        max_ulysses_dim_size = int(world_size / max_batch_dim_size)
+        logger.info(f"max_batch_dim_size: {max_batch_dim_size}, max_ulysses_dim_size:{max_ulysses_dim_size}")
+
+        mesh = init_context_parallel_mesh(
+            self.pipe.device.type,
+            max_ring_dim_size=1,
+            max_batch_dim_size=max_batch_dim_size,
+        )
+        parallelize_pipe(self.pipe, mesh=mesh)
+        parallelize_vae(self.pipe.vae, mesh=mesh._flatten())
+
+        if is_offload:
+            Offload.offload(
+                pipeline=self.pipe,
+                config=offload_config,
+            )
+        else:
+            self.pipe.to(gpu_device)
+
+        if offload_config.compiler_transformer:
+            torch._dynamo.config.suppress_errors = True
+            os.environ["TORCHINDUCTOR_FX_GRAPH_CACHE"] = "1"
+            os.environ["TORCHINDUCTOR_CACHE_DIR"] = f"{offload_config.compiler_cache}_{world_size}"
+            self.pipe.transformer = torch.compile(
+                self.pipe.transformer,
+                mode="max-autotune-no-cudagraphs",
+                dynamic=True,
+            )
+            self.warm_up()
+
+    def warm_up(self):
+        init_kwargs = {
+            "prompt": "A woman is dancing in a room",
+            "height": 544,
+            "width": 960,
+            "guidance_scale": 6,
+            "num_inference_steps": 1,
+            "negative_prompt": "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion",
+            "num_frames": 97,
+            "generator": torch.Generator("cuda").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)
+
+    def damon_inference(self, request_queue: mp.Queue, response_queue: mp.Queue):
+        response_queue.put(f"rank:{self.gpu_rank} ready")
+        logger.info(f"rank:{self.gpu_rank} finish init pipe")
+        while True:
+            logger.info(f"rank:{self.gpu_rank} waiting for request")
+            kwargs = request_queue.get()
+            logger.info(f"rank:{self.gpu_rank} kwargs: {kwargs}")
+            if "seed" in kwargs:
+                kwargs["generator"] = torch.Generator("cuda").manual_seed(kwargs["seed"])
+                del kwargs["seed"]
+            start_time = time.time()
+            assert (self.task_type == TaskType.I2V and "image" in kwargs) or self.task_type == TaskType.T2V
+            out = self.pipe(**kwargs).frames[0]
+            logger.info(f"rank:{dist.get_rank()} inference time: {time.time() - start_time}")
+            if dist.get_rank() == 0:
+                response_queue.put(out)
+
+
+def single_gpu_run(
+    rank,
+    task_type: TaskType,
+    model_id: str,
+    request_queue: mp.Queue,
+    response_queue: mp.Queue,
+    quant_model: bool = True,
+    world_size: int = 1,
+    is_offload: bool = True,
+    offload_config: OffloadConfig = OffloadConfig(),
+    enable_cfg_parallel: bool = True,
+):
+    pipe = SkyReelsVideoSingleGpuInfer(
+        task_type=task_type,
+        model_id=model_id,
+        quant_model=quant_model,
+        local_rank=rank,
+        world_size=world_size,
+        is_offload=is_offload,
+        offload_config=offload_config,
+        enable_cfg_parallel=enable_cfg_parallel,
+    )
+    pipe.damon_inference(request_queue, response_queue)
+
+
+class SkyReelsVideoInfer:
+    def __init__(
+        self,
+        task_type: TaskType,
+        model_id: str,
+        quant_model: bool = True,
+        world_size: int = 1,
+        is_offload: bool = True,
+        offload_config: OffloadConfig = OffloadConfig(),
+        enable_cfg_parallel: bool = True,
+    ):
+        self.world_size = world_size
+        smp = mp.get_context("spawn")
+        self.REQ_QUEUES: mp.Queue = smp.Queue()
+        self.RESP_QUEUE: mp.Queue = smp.Queue()
+        assert self.world_size > 0, "gpu_num must be greater than 0"
+        spawn_thread = threading.Thread(
+            target=self.lauch_single_gpu_infer,
+            args=(task_type, model_id, quant_model, world_size, is_offload, offload_config, enable_cfg_parallel),
+            daemon=True,
+        )
+        spawn_thread.start()
+        logger.info(f"Started multi-GPU thread with GPU_NUM: {world_size}")
+        print(f"Started multi-GPU thread with GPU_NUM: {world_size}")
+        # Block and wait for the prediction process to start
+        for _ in range(world_size):
+            msg = self.RESP_QUEUE.get()
+            logger.info(f"launch_multi_gpu get init msg: {msg}")
+            print(f"launch_multi_gpu get init msg: {msg}")
+
+    def lauch_single_gpu_infer(
+        self,
+        task_type: TaskType,
+        model_id: str,
+        quant_model: bool = True,
+        world_size: int = 1,
+        is_offload: bool = True,
+        offload_config: OffloadConfig = OffloadConfig(),
+        enable_cfg_parallel: bool = True,
+    ):
+        mp.spawn(
+            single_gpu_run,
+            nprocs=world_size,
+            join=True,
+            daemon=True,
+            args=(
+                task_type,
+                model_id,
+                self.REQ_QUEUES,
+                self.RESP_QUEUE,
+                quant_model,
+                world_size,
+                is_offload,
+                offload_config,
+                enable_cfg_parallel,
+            ),
+        )
+        logger.info(f"finish lanch multi gpu infer, world_size:{world_size}")
+
+    def inference(self, kwargs: Dict[str, Any]):
+        # put request to singlegpuinfer
+        for _ in range(self.world_size):
+            self.REQ_QUEUES.put(kwargs)
+        return self.RESP_QUEUE.get()