xfuser 0.1__py3-none-any.whl

pipefuser/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.1"

pipefuser/logger.py

@@ -0,0 +1,99 @@
+# Adapted from
+# https://github.com/skypilot-org/skypilot/blob/86dc0f6283a335e4aa37b3c10716f90999f48ab6/sky/sky_logging.py
+"""Logging configuration."""
+import logging
+import sys
+import os
+from typing import Optional
+
+_FORMAT = "%(levelname)s %(asctime)s [%(filename)s:%(lineno)d] %(message)s"
+_DATE_FORMAT = "%m-%d %H:%M:%S"
+
+_LOG_LEVEL = os.environ.get("LOG_LEVEL", "debug")
+_LOG_LEVEL = getattr(logging, _LOG_LEVEL.upper(), 0)
+_LOG_DIR = os.environ.get("LOG_DIR", None)
+
+
+class NewLineFormatter(logging.Formatter):
+    """Adds logging prefix to newlines to align multi-line messages."""
+
+    def __init__(self, fmt, datefmt=None):
+        logging.Formatter.__init__(self, fmt, datefmt)
+
+    def format(self, record):
+        msg = logging.Formatter.format(self, record)
+        if record.message != "":
+            parts = msg.split(record.message)
+            msg = msg.replace("\n", "\r\n" + parts[0])
+        return msg
+
+
+_root_logger = logging.getLogger("pipefuser")
+_default_handler = None
+_default_file_handler = None
+_inference_log_file_handler = {}
+
+
+def _setup_logger():
+    _root_logger.setLevel(_LOG_LEVEL)
+    global _default_handler
+    global _default_file_handler
+    fmt = NewLineFormatter(_FORMAT, datefmt=_DATE_FORMAT)
+
+    if _default_handler is None:
+        _default_handler = logging.StreamHandler(sys.stdout)
+        _default_handler.flush = sys.stdout.flush  # type: ignore
+        _default_handler.setLevel(_LOG_LEVEL)
+        _root_logger.addHandler(_default_handler)
+
+    if _default_file_handler is None and _LOG_DIR is not None:
+        if not os.path.exists(_LOG_DIR):
+            try:
+                os.makedirs(_LOG_DIR)
+            except OSError as e:
+                _root_logger.warn(f"Error creating directory {_LOG_DIR} : {e}")
+        _default_file_handler = logging.FileHandler(_LOG_DIR + "/default.log")
+        _default_file_handler.setLevel(_LOG_LEVEL)
+        _default_file_handler.setFormatter(fmt)
+        _root_logger.addHandler(_default_file_handler)
+
+    _default_handler.setFormatter(fmt)
+    # Setting this will avoid the message
+    # being propagated to the parent logger.
+    _root_logger.propagate = False
+
+
+# The logger is initialized when the module is imported.
+# This is thread-safe as the module is only imported once,
+# guaranteed by the Python GIL.
+_setup_logger()
+
+
+def init_logger(name: str):
+    pid = os.getpid()
+    # Use the same settings as above for root logger
+    logger = logging.getLogger(name)
+    logger.setLevel(_LOG_LEVEL)
+    logger.addHandler(_default_handler)
+    if _LOG_DIR is not None and pid is None:
+        logger.addHandler(_default_file_handler)
+    elif _LOG_DIR is not None:
+        if _inference_log_file_handler.get(pid, None) is not None:
+            logger.addHandler(_inference_log_file_handler[pid])
+        else:
+            if not os.path.exists(_LOG_DIR):
+                try:
+                    os.makedirs(_LOG_DIR)
+                except OSError as e:
+                    _root_logger.warn(f"Error creating directory {_LOG_DIR} : {e}")
+            _inference_log_file_handler[pid] = logging.FileHandler(
+                _LOG_DIR + f"/process.{pid}.log"
+            )
+            _inference_log_file_handler[pid].setLevel(_LOG_LEVEL)
+            _inference_log_file_handler[pid].setFormatter(
+                NewLineFormatter(_FORMAT, datefmt=_DATE_FORMAT)
+            )
+            _root_logger.addHandler(_inference_log_file_handler[pid])
+            logger.addHandler(_inference_log_file_handler[pid])
+    logger.propagate = False
+    return logger

pipefuser/models/__init__.py

@@ -0,0 +1,13 @@
+from packaging.version import Version
+import diffusers
+
+if Version('0.29.0') <= Version(diffusers.__version__):
+    from .distri_dit_sd3_pipefusion import DistriDiTSD3PipeFusion
+
+from .naive_patch_dit import NaivePatchDiT
+from .distri_dit_pp import DistriDiTPP
+from .distri_dit_tp import DistriDiTTP
+from .distri_dit_pipefusion import DistriDiTPipeFusion
+from .distri_dit_hunyuan_pipefusion import DistriDiTHunyuanPipeFusion
+
+

pipefuser/models/base_model.py

@@ -0,0 +1,52 @@
+from diffusers import ConfigMixin, ModelMixin
+from torch import nn
+
+from pipefuser.modules.base_module import BaseModule
+from ..utils import PatchParallelismCommManager, DistriConfig
+
+
+class BaseModel(ModelMixin, ConfigMixin):
+    def __init__(self, model: nn.Module, distri_config: DistriConfig):
+        super(BaseModel, self).__init__()
+        self.model = model
+        self.distri_config = distri_config
+        self.comm_manager = None
+
+        self.buffer_list = None
+        self.output_buffer = None
+        self.counter = 0
+
+        # for cuda graph
+        self.static_inputs = None
+        self.static_outputs = None
+        self.cuda_graphs = None
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def set_counter(self, counter: int = 0):
+        self.counter = counter
+        for module in self.model.modules():
+            if isinstance(module, BaseModule):
+                module.set_counter(counter)
+
+    def set_comm_manager(self, comm_manager: PatchParallelismCommManager):
+        self.comm_manager = comm_manager
+        for module in self.model.modules():
+            if isinstance(module, BaseModule):
+                module.set_comm_manager(comm_manager)
+
+    def setup_cuda_graph(self, static_outputs, cuda_graphs):
+        self.static_outputs = static_outputs
+        self.cuda_graphs = cuda_graphs
+
+    @property
+    def config(self):
+        return self.model.config
+
+    def synchronize(self):
+        if self.comm_manager is not None and self.comm_manager.handles is not None:
+            for i in range(len(self.comm_manager.handles)):
+                if self.comm_manager.handles[i] is not None:
+                    self.comm_manager.handles[i].wait()
+                    self.comm_manager.handles[i] = None

pipefuser/models/diffusers/__init__.py

@@ -0,0 +1,10 @@
+USE_CUSTOMIZED_IMPL = False
+
+if USE_CUSTOMIZED_IMPL:
+    from .attention import Attention
+    from .transformers_2d import Transformer2DModel
+else:
+    from diffusers.models.attention import Attention
+    from diffusers.models.transformers.transformer_2d import Transformer2DModel
+
+__all__ = ["Attention", "Transformer2DModel"]