diffsynth-engine 0.0.0__py3-none-any.whl

diffsynth_engine/utils/lock.py

@@ -0,0 +1,56 @@
+import math
+import threading
+from typing import Optional
+from types import TracebackType
+from flufl.lock import Lock
+
+from diffsynth_engine.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class HeartbeatFileLock:
+    def __init__(self, lock_file_path: str, heartbeat_interval: float = 10):
+        self.lock_file_path = lock_file_path
+        self.heartbeat_interval = heartbeat_interval
+        self.lifetime = math.ceil(heartbeat_interval + 1)
+        self.heartbeat_thread = None
+        self.stop_event = threading.Event()
+        self.lock = None
+
+    def _heartbeat(self):
+        while not self.stop_event.is_set():
+            self.lock.refresh(lifetime=self.lifetime)
+            self.stop_event.wait(self.heartbeat_interval - 1)
+
+    def acquire(self):
+        self.lock = Lock(self.lock_file_path, lifetime=self.lifetime)
+        self.lock.lock()
+
+        self.heartbeat_thread = threading.Thread(target=self._heartbeat)
+        self.heartbeat_thread.start()
+
+    def release(self):
+        if self.lock is not None:
+            self.lock.unlock(unconditionally=True)
+        self._release()
+
+    def _release(self):
+        if self.heartbeat_thread is not None:
+            self.stop_event.set()
+            self.heartbeat_thread.join()
+
+    def __enter__(self):
+        self.acquire()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]] = None,
+        exc_value: Optional[BaseException] = None,
+        traceback: Optional[TracebackType] = None,
+    ):
+        self._release()
+
+    def __del__(self):
+        self.release()

diffsynth_engine/utils/logging.py

@@ -0,0 +1,12 @@
+import logging
+from typing import Optional
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+
+
+def get_logger(name: Optional[str] = None) -> logging.Logger:
+    return logging.getLogger(name)

diffsynth_engine/utils/offload.py

@@ -0,0 +1,44 @@
+import torch.nn as nn
+
+from diffsynth_engine.models.basic.transformer_helper import RMSNorm
+from diffsynth_engine.models.basic.relative_position_emb import RelativePositionEmbedding
+
+
+SUPPORTED_OFFLOAD_MODULES = (
+    nn.Embedding,
+    nn.Linear,
+    nn.LayerNorm,
+    nn.Conv2d,
+    nn.GroupNorm,
+    RMSNorm,
+    RelativePositionEmbedding,
+)
+
+
+def enable_sequential_cpu_offload(module: nn.Module, device: str = "cuda:0"):
+    if isinstance(module, SUPPORTED_OFFLOAD_MODULES):
+        add_cpu_offload_hook(module, device)
+        return
+    for submodule in module.children():
+        enable_sequential_cpu_offload(submodule, device)
+
+
+def add_cpu_offload_hook(module: nn.Module, device: str = "cuda:0"):
+    def _forward_pre_hook(module: nn.Module, input):
+        offload_params = {}
+        for name, param in module.named_parameters():
+            offload_params[name] = param.data
+            param.data = param.data.to(device=device)
+        setattr(module, "_offload_params", offload_params)
+
+    def _forward_hook(module: nn.Module, input, output):
+        offload_params = getattr(module, "_offload_params", {})
+        for name, param in module.named_parameters():
+            if name in offload_params:
+                param.data = offload_params[name]
+
+    if getattr(module, "_sequential_cpu_offload_enabled", False):
+        return
+    module.register_forward_pre_hook(_forward_pre_hook)
+    module.register_forward_hook(_forward_hook)
+    setattr(module, "_sequential_cpu_offload_enabled", True)

diffsynth_engine/utils/parallel.py

@@ -0,0 +1,390 @@
+import os
+import copy
+import torch
+import torch.nn as nn
+import torch.multiprocessing as mp
+import torch.distributed as dist
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import ShardingStrategy
+from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor.parallel.style import ParallelStyle
+from torch.distributed.tensor.parallel._utils import _validate_tp_mesh_dim
+from datetime import timedelta
+from functools import partial
+from typing import Callable, Dict, List, Union, Optional
+from queue import Empty
+from yunchang.globals import Singleton, set_seq_parallel_pg
+
+from diffsynth_engine.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class ProcessGroupSingleton(Singleton):
+    def __init__(self):
+        self.CFG_GROUP: dist.ProcessGroup = None
+        self.SP_GROUP: dist.ProcessGroup = None
+        self.TP_GROUP: dist.ProcessGroup = None
+
+        self.CFG_RANKS: List[int] = []
+        self.SP_RANKS: List[int] = []
+        self.TP_RANKS: List[int] = []
+
+
+PROCESS_GROUP = ProcessGroupSingleton()
+
+
+def get_cfg_group():
+    return PROCESS_GROUP.CFG_GROUP
+
+
+def get_cfg_world_size():
+    return PROCESS_GROUP.CFG_GROUP.size()
+
+
+def get_cfg_rank():
+    return PROCESS_GROUP.CFG_GROUP.rank()
+
+
+def get_cfg_ranks():
+    return PROCESS_GROUP.CFG_RANKS
+
+
+def get_sp_group():
+    return PROCESS_GROUP.SP_GROUP
+
+
+def get_sp_world_size():
+    return PROCESS_GROUP.SP_GROUP.size()
+
+
+def get_sp_rank():
+    return PROCESS_GROUP.SP_GROUP.rank()
+
+
+def get_sp_ranks():
+    return PROCESS_GROUP.SP_RANKS
+
+
+def get_tp_group():
+    return PROCESS_GROUP.TP_GROUP
+
+
+def get_tp_world_size():
+    return PROCESS_GROUP.TP_GROUP.size()
+
+
+def get_tp_rank():
+    return PROCESS_GROUP.TP_GROUP.rank()
+
+
+def get_tp_ranks():
+    return PROCESS_GROUP.TP_RANKS
+
+
+def init_parallel_pgs(
+    cfg_degree: int = 1,
+    sp_ulysses_degree: int = 1,
+    sp_ring_degree: int = 1,
+    tp_degree: int = 1,
+    rank: int = 0,
+    world_size: int = 1,
+):
+    sp_degree = sp_ulysses_degree * sp_ring_degree
+
+    assert sp_degree == 1 or tp_degree == 1, "not allowed to enable sequence parallel and tensor parallel together"
+    assert world_size == cfg_degree * sp_degree * tp_degree, (
+        f"world_size ({world_size}) must be equal to cfg_degree ({cfg_degree}) * sp_degree ({sp_degree}) * tp_degree ({tp_degree})"
+    )
+
+    def make_parallel_groups(blocks: List[List[int]], degree: int):
+        groups, chunks = [], []
+        for block in blocks:
+            size = len(block) // degree
+            chunk = [block[i * size : (i + 1) * size] for i in range(degree)]
+            chunks.extend(chunk)
+            groups.extend(list(zip(*chunk)))
+        return groups, chunks
+
+    blocks = [list(range(world_size))]
+    cfg_groups, cfg_blocks = make_parallel_groups(blocks, cfg_degree)
+    for cfg_ranks in cfg_groups:
+        cfg_group = dist.new_group(cfg_ranks)
+        if rank in cfg_ranks:
+            PROCESS_GROUP.CFG_GROUP = cfg_group
+            PROCESS_GROUP.CFG_RANKS = cfg_ranks
+
+    sp_groups, sp_blocks = make_parallel_groups(cfg_blocks, sp_degree)
+    for sp_ranks in sp_groups:
+        group = dist.new_group(sp_ranks)
+        if rank in sp_ranks:
+            PROCESS_GROUP.SP_GROUP = group
+            PROCESS_GROUP.SP_RANKS = sp_ranks
+
+    tp_groups, _ = make_parallel_groups(sp_blocks, tp_degree)
+    for tp_ranks in tp_groups:
+        group = dist.new_group(tp_ranks)
+        if rank in tp_ranks:
+            PROCESS_GROUP.TP_GROUP = group
+            PROCESS_GROUP.TP_RANKS = tp_ranks
+
+    set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size)
+
+
+def clone(data):
+    if isinstance(data, dict):
+        return {k: clone(v) for k, v in data.items()}
+    if isinstance(data, tuple) or isinstance(data, list):
+        return [clone(t) for t in data]
+    elif isinstance(data, torch.Tensor):
+        return data.clone()
+    else:
+        return copy.deepcopy(data)
+
+
+def to_device(data, device):
+    if isinstance(data, dict):
+        return {k: to_device(v, device) for k, v in data.items()}
+    if isinstance(data, tuple) or isinstance(data, list):
+        return [to_device(t, device) for t in data]
+    elif isinstance(data, torch.Tensor):
+        return data.to(device)
+    else:
+        return data
+
+
+def split_and_get(data, num, dim, index):
+    if isinstance(data, dict):
+        return {k: split_and_get(v, num, dim, index) for k, v in data.items()}
+    if isinstance(data, tuple) or isinstance(data, list):
+        return [split_and_get(t, num, dim, index) for t in data]
+    if isinstance(data, torch.Tensor):
+        if data.shape[dim] < num:
+            raise ValueError(f"data.shape[{dim}] ({data.shape[dim]}) < num ({num}), split failed")
+        return torch.split(data, data.shape[dim] // num, dim)[index]
+    return data
+
+
+def shard_model(
+    module: nn.Module,
+    device_id: int,
+    sharding_strategy: ShardingStrategy = ShardingStrategy.FULL_SHARD,
+    wrap_module_names: Optional[List[str]] = None,
+):
+    wrap_module_names = wrap_module_names or []
+
+    def wrap_fn(m):
+        for name in wrap_module_names:
+            submodule = getattr(module, name)
+            if isinstance(submodule, nn.ModuleList) and m in submodule:
+                return True
+            elif not isinstance(submodule, nn.ModuleList) and m is submodule:
+                return True
+        return False
+
+    return FSDP(
+        module,
+        device_id=device_id,
+        sharding_strategy=sharding_strategy,
+        auto_wrap_policy=partial(lambda_auto_wrap_policy, lambda_fn=wrap_fn),
+    )
+
+
+def parallelize_module(
+    module: nn.Module,
+    device_mesh: DeviceMesh,
+    parallelize_plan: Optional[Union[ParallelStyle, Dict[str, ParallelStyle]]] = None,
+):
+    _validate_tp_mesh_dim(device_mesh)
+    if parallelize_plan is None:
+        return module
+    if isinstance(parallelize_plan, ParallelStyle):
+        return parallelize_plan._apply(module, device_mesh)
+    for module_path, parallelize_style in parallelize_plan.items():
+        if module_path.strip() == "":
+            raise ValueError("Expect module path to be non-empty, but got empty string!")
+        try:
+            submodule = module.get_submodule(module_path)
+            parallelize_style._apply(submodule, device_mesh)
+        except AttributeError:
+            continue
+    return module
+
+
+NCCL_TIMEOUT_SEC = int(os.environ.get("NCCL_TIMEOUT_SEC", 600))
+PARALLEL_FWD_TIMEOUT_SEC = int(os.environ.get("PARALLEL_FWD_TIMEOUT_SEC", 300))
+PARALLEL_LORA_TIMEOUT_SEC = int(os.environ.get("PARALLEL_LORA_TIMEOUT_SEC ", 60))
+
+
+def _worker_loop(
+    rank: int,
+    world_size: int,
+    queue_in: mp.Queue,
+    queue_out: mp.Queue,
+    module: nn.Module,
+    cfg_degree: int,
+    sp_ulysses_degree: int,
+    sp_ring_degree: int,
+    tp_degree: int,
+    shard_fn: Optional[Callable] = None,
+    master_port: int = 29500,
+    device: str = "cuda",
+):
+    """
+    https://pytorch.org/docs/stable/multiprocessing.html#sharing-cuda-tensors
+    """
+    try:
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(master_port)
+        torch.cuda.set_device(rank)
+
+        timeout = timedelta(seconds=NCCL_TIMEOUT_SEC)
+        dist.init_process_group(
+            backend="nccl",
+            init_method="env://",
+            timeout=timeout,
+            world_size=world_size,
+            rank=rank,
+        )
+        init_parallel_pgs(
+            cfg_degree=cfg_degree,
+            sp_ulysses_degree=sp_ulysses_degree,
+            sp_ring_degree=sp_ring_degree,
+            tp_degree=tp_degree,
+            rank=rank,
+            world_size=world_size,
+        )
+
+        if tp_degree > 1:
+            module = parallelize_module(
+                module=module,
+                device_mesh=DeviceMesh(device, torch.tensor(get_tp_ranks())),
+                parallelize_plan=module.get_tp_plan(),
+            ).to(device)
+        elif shard_fn:
+            module = shard_fn(module=module, device_id=rank)
+        else:
+            module = module.to(device)
+
+        while True:
+            if rank == 0:
+                kwargs = queue_in.get()
+                data = [kwargs]
+            else:
+                data = [None]
+            dist.broadcast_object_list(data, src=0)
+            kwargs = clone(data[0])
+            del data
+
+            y = None
+            if kwargs.get("method", None) == "load_loras":
+                module.load_loras(lora_args=kwargs["lora_args"], fused=kwargs["fused"])
+            elif kwargs.get("method", None) == "unload_loras":
+                module.unload_loras()
+            else:
+                kwargs = to_device(kwargs, device)
+                kwargs = split_and_get(kwargs, get_cfg_world_size(), 0, get_cfg_rank())
+                with torch.no_grad():
+                    y = module(**kwargs)
+                if get_sp_rank() == 0 and get_tp_rank() == 0:
+                    gathered = torch.zeros((get_cfg_world_size(), *y.shape[1:]), dtype=y.dtype, device=y.device)
+                    dist.all_gather_into_tensor(gathered, y, group=get_cfg_group())
+                    y = gathered
+
+            if rank == 0:
+                queue_out.put(y)
+            dist.barrier()
+    except Exception as e:
+        import traceback
+
+        traceback.print_exc()
+        logger.error(f"Error in worker loop (rank {rank}): {e}")
+    finally:
+        del module
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+        dist.destroy_process_group()
+
+
+class ParallelModel(nn.Module):
+    def __init__(
+        self,
+        module: nn.Module,
+        cfg_degree: int,
+        sp_ulysses_degree: int,
+        sp_ring_degree: int,
+        tp_degree: int,
+        shard_fn: Optional[Callable] = None,
+        master_port: int = 29500,
+        device: str = "cuda",
+    ):
+        super().__init__()
+        self.world_size = cfg_degree * sp_ulysses_degree * sp_ring_degree * tp_degree
+        self.device = device
+        self.queue_in = mp.Queue()
+        self.queue_out = mp.Queue()
+        self.ctx = mp.spawn(
+            _worker_loop,
+            args=(
+                self.world_size,
+                self.queue_in,
+                self.queue_out,
+                module,
+                cfg_degree,
+                sp_ulysses_degree,
+                sp_ring_degree,
+                tp_degree,
+                shard_fn,
+                master_port,
+                device,
+            ),
+            nprocs=self.world_size,
+            join=False,
+        )
+
+    def load_loras(self, lora_args: List[Dict[str, any]], fused: bool = True):
+        self.queue_in.put(
+            {
+                "method": "load_loras",
+                "lora_args": lora_args,
+                "fused": fused,
+            }
+        )
+        try:
+            _ = self.queue_out.get(timeout=PARALLEL_LORA_TIMEOUT_SEC)
+        except Empty:
+            logger.error("Parallel model load LoRA timeout")
+            raise RuntimeError("Parallel model load LoRA timeout")
+        logger.info("Parallel model load LoRA done")
+
+    def unload_loras(self):
+        self.queue_in.put({"method": "unload_loras"})
+        try:
+            _ = self.queue_out.get(timeout=PARALLEL_LORA_TIMEOUT_SEC)
+        except Empty:
+            logger.error("Parallel model unload LoRA timeout")
+            raise RuntimeError("Parallel model unload LoRA timeout")
+        logger.info("Parallel model unload LoRA done")
+
+    def forward(self, **kwargs):
+        self.queue_in.put(kwargs)
+        try:
+            y = self.queue_out.get(timeout=PARALLEL_FWD_TIMEOUT_SEC)
+        except Empty:
+            logger.error("Parallel model forward timeout")
+            raise RuntimeError("Parallel model forward timeout")
+        return y
+
+    def __del__(self):
+        # Send terminate signal to all workers
+        for p in self.ctx.processes:
+            p.terminate()
+            p.join()
+        self.queue_in.close()
+        self.queue_out.close()
+
+
+__all__ = ["ParallelModel"]

diffsynth_engine/utils/prompt.py

@@ -0,0 +1,9 @@
+from typing import Union, Optional
+
+from diffsynth_engine.tokenizers import CLIPTokenizer, T5TokenizerFast
+
+
+def tokenize_long_prompt(
+    tokenizer: Union[CLIPTokenizer, T5TokenizerFast], prompt: str, max_length: Optional[int] = None
+):
+    return tokenizer(prompt)["input_ids"]

diffsynth_engine/utils/video.py

@@ -0,0 +1,40 @@
+import imageio
+import imageio.v3 as iio
+import numpy as np
+from PIL import Image
+from typing import List
+
+
+class VideoReader:
+    def __init__(self, path: str):
+        self.reader = imageio.get_reader(path)
+
+    def __len__(self):
+        return self.reader.count_frames()
+
+    def __getitem__(self, item):
+        return Image.fromarray(np.array(self.reader.get_data(item))).convert("RGB")
+
+    def __del__(self):
+        self.reader.close()
+
+    @property
+    def frames(self) -> List[Image.Image]:
+        return [self[i] for i in range(len(self))]
+
+
+def load_video(path: str) -> VideoReader:
+    return VideoReader(path)
+
+
+def save_video(frames, save_path, fps=15):
+    if save_path.endswith(".webm"):
+        codec = "libvpx-vp9"
+    elif save_path.endswith(".mp4"):
+        codec = "libx264"
+
+    frames = [np.array(img) for img in frames]
+
+    # 使用 imageio 写入 .webm 文件
+    with iio.imopen(save_path, "w", plugin="FFMPEG") as writer:
+        writer.write(frames, fps=fps, codec=codec)