chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/utils/camera_geometry.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Camera pose and Plucker ray utilities."""
+
+from __future__ import annotations
+
+import torch
+
+
+def se3_inverse(T: torch.Tensor) -> torch.Tensor:
+    rot = T[:, :3, :3]
+    trans = T[:, :3, 3:]
+    r_inv = rot.transpose(-1, -2)
+    t_inv = -torch.bmm(r_inv, trans)
+    T_inv = torch.eye(4, device=T.device, dtype=T.dtype)[None, :, :].repeat(
+        T.shape[0], 1, 1
+    )
+    T_inv[:, :3, :3] = r_inv
+    T_inv[:, :3, 3:] = t_inv
+    return T_inv
+
+
+def compute_relative_poses(
+    c2ws_mat: torch.Tensor,
+    framewise: bool = False,
+    normalize_trans: bool = True,
+) -> torch.Tensor:
+    ref_w2cs = se3_inverse(c2ws_mat[0:1])
+    relative_poses = torch.matmul(ref_w2cs, c2ws_mat)
+    relative_poses[0] = torch.eye(4, device=c2ws_mat.device, dtype=c2ws_mat.dtype)
+    if framewise and len(relative_poses) > 1:
+        relative_poses_framewise = torch.bmm(
+            se3_inverse(relative_poses[:-1]), relative_poses[1:]
+        )
+        relative_poses[1:] = relative_poses_framewise
+    if normalize_trans:
+        translations = relative_poses[:, :3, 3]
+        max_norm = torch.norm(translations, dim=-1).max()
+        if max_norm > 0:
+            relative_poses[:, :3, 3] = translations / max_norm
+    return relative_poses
+
+
+@torch.no_grad()
+def create_meshgrid(
+    n_frames: int,
+    height: int,
+    width: int,
+    *,
+    bias: float = 0.5,
+    device: torch.device | str,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    x_range = torch.arange(width, device=device, dtype=dtype)
+    y_range = torch.arange(height, device=device, dtype=dtype)
+    grid_y, grid_x = torch.meshgrid(y_range, x_range, indexing="ij")
+    grid_xy = torch.stack([grid_x, grid_y], dim=-1).view([-1, 2]) + bias
+    return grid_xy[None, ...].repeat(n_frames, 1, 1)
+
+
+def get_plucker_embeddings(
+    c2ws_mat: torch.Tensor,
+    Ks: torch.Tensor,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    n_frames = c2ws_mat.shape[0]
+    grid_xy = create_meshgrid(
+        n_frames, height, width, device=c2ws_mat.device, dtype=c2ws_mat.dtype
+    )
+    fx, fy, cx, cy = Ks.chunk(4, dim=-1)
+    i = grid_xy[..., 0]
+    j = grid_xy[..., 1]
+    zs = torch.ones_like(i)
+    xs = (i - cx) / fx * zs
+    ys = (j - cy) / fy * zs
+    directions = torch.stack([xs, ys, zs], dim=-1)
+    directions = directions / directions.norm(dim=-1, keepdim=True)
+    rays_d = directions @ c2ws_mat[:, :3, :3].transpose(-1, -2)
+    rays_o = c2ws_mat[:, :3, 3][:, None, :].expand_as(rays_d)
+    plucker_embeddings = torch.cat([rays_o, rays_d], dim=-1)
+    return plucker_embeddings.view([n_frames, height, width, 6])
+
+
+def camera_poses_to_plucker(
+    *,
+    c2ws: torch.Tensor,
+    Ks: torch.Tensor,
+    height: int,
+    width: int,
+    spatial_scale: int = 8,
+    device: torch.device | str,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    plucker = get_plucker_embeddings(c2ws, Ks, height, width)
+    latent_height = height // spatial_scale
+    latent_width = width // spatial_scale
+    plucker = plucker.view(
+        c2ws.shape[0],
+        latent_height,
+        spatial_scale,
+        latent_width,
+        spatial_scale,
+        6,
+    )
+    plucker = plucker.permute(0, 1, 3, 5, 2, 4).contiguous()
+    plucker = plucker.view(
+        c2ws.shape[0],
+        latent_height,
+        latent_width,
+        6 * spatial_scale * spatial_scale,
+    )
+    return (
+        plucker.permute(3, 0, 1, 2)
+        .contiguous()
+        .unsqueeze(0)
+        .to(device=device, dtype=dtype)
+    )

python/sglang/multimodal_gen/runtime/utils/common.py

@@ -0,0 +1,423 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import ipaddress
+import logging
+import os
+import platform
+import signal
+import socket
+import sys
+import threading
+from functools import lru_cache
+from typing import Any
+
+import psutil
+import torch
+import zmq
+
+# use the native logger to avoid circular import
+logger = logging.getLogger(__name__)
+
+
+def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
+    """Kill the process and all its child processes."""
+    # Remove sigchld handler to avoid spammy logs.
+    if threading.current_thread() is threading.main_thread():
+        signal.signal(signal.SIGCHLD, signal.SIG_DFL)
+
+    if parent_pid is None:
+        parent_pid = os.getpid()
+        include_parent = False
+
+    try:
+        itself = psutil.Process(parent_pid)
+    except psutil.NoSuchProcess:
+        return
+
+    children = itself.children(recursive=True)
+    for child in children:
+        if child.pid == skip_pid:
+            continue
+        try:
+            child.kill()
+        except psutil.NoSuchProcess:
+            pass
+
+    if include_parent:
+        try:
+            if parent_pid == os.getpid():
+                itself.kill()
+                sys.exit(0)
+
+            itself.kill()
+
+            # Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes),
+            # so we send an additional signal to kill them.
+            itself.send_signal(signal.SIGQUIT)
+        except psutil.NoSuchProcess:
+            pass
+
+
+def add_prefix(name: str, prefix: str) -> str:
+    """Add a weight path prefix to a module name.
+
+    Args:
+        name: base module name.
+        prefix: weight prefix str to added to the front of `name` concatenated with `.`.
+
+    Returns:
+        The string `prefix.name` if prefix is non-empty, otherwise just `name`.
+    """
+    return name if not prefix else f"{prefix}.{name}"
+
+
+def is_valid_ipv6_address(address: str) -> bool:
+    try:
+        ipaddress.IPv6Address(address)
+        return True
+    except ValueError:
+        return False
+
+
+def normalize_gpu_ids(gpu_ids: Any) -> list[int] | None:
+    if gpu_ids is None:
+        return None
+    if isinstance(gpu_ids, str):
+        values = [gpu_ids]
+    else:
+        values = list(gpu_ids)
+
+    tokens: list[str] = []
+    for value in values:
+        tokens.extend(part for part in str(value).replace(",", " ").split() if part)
+    if not tokens:
+        return []
+
+    parsed: list[int] = []
+    for token in tokens:
+        try:
+            gpu_id = int(token)
+        except ValueError as exc:
+            raise ValueError(
+                f"--gpu-ids contains a non-integer GPU id: {token}"
+            ) from exc
+        if gpu_id < 0:
+            raise ValueError(f"--gpu-ids GPU ids must be non-negative: {gpu_id}")
+        parsed.append(gpu_id)
+
+    if len(set(parsed)) != len(parsed):
+        raise ValueError(f"--gpu-ids contains duplicate GPU ids: {parsed}")
+    return parsed
+
+
+def parse_size(size: str) -> tuple[int | None, int | None]:
+    try:
+        parts = size.lower().replace(" ", "").split("x")
+        if len(parts) != 2:
+            raise ValueError
+        return int(parts[0]), int(parts[1])
+    except ValueError:
+        return None, None
+
+
+def parse_tcp_host_port(value: str | None, field_name: str) -> tuple[str, int]:
+    if value is None or not str(value).strip():
+        raise ValueError(f"{field_name} is required")
+
+    addr = str(value).strip()
+    if addr.startswith("tcp://"):
+        addr = addr[len("tcp://") :]
+
+    try:
+        host, port_str = addr.rsplit(":", 1)
+    except ValueError as exc:
+        raise ValueError(
+            f"{field_name} must be formatted as tcp://host:port or host:port"
+        ) from exc
+
+    host = host.strip()
+    port_str = port_str.strip()
+    if not host or not port_str:
+        raise ValueError(f"{field_name} must include both host and port: {value!r}")
+
+    try:
+        port = int(port_str)
+    except ValueError as exc:
+        raise ValueError(f"{field_name} port must be an integer: {port_str}") from exc
+
+    if port < 0 or port > 65535:
+        raise ValueError(f"{field_name} port must be between 0 and 65535: {port}")
+    return host, port
+
+
+def format_tcp_endpoint(host: str, port: int, field_name: str) -> str:
+    if port < 0 or port > 65535:
+        raise ValueError(f"{field_name} port must be between 0 and 65535: {port}")
+    return f"tcp://{host}:{port}"
+
+
+def configure_ipv6(dist_init_addr):
+    addr = dist_init_addr
+    end = addr.find("]")
+    if end == -1:
+        raise ValueError("invalid IPv6 address format: missing ']'")
+
+    host = addr[: end + 1]
+
+    # this only validates the address without brackets: we still need the below checks.
+    # if it's invalid, immediately raise an error so we know it's not formatting issues.
+    if not is_valid_ipv6_address(host[1:end]):
+        raise ValueError(f"invalid IPv6 address: {host}")
+
+    port_str = None
+    if len(addr) > end + 1:
+        if addr[end + 1] == ":":
+            port_str = addr[end + 2 :]
+        else:
+            raise ValueError("received IPv6 address format: expected ':' after ']'")
+
+    if not port_str:
+        raise ValueError(
+            "a port must be specified in IPv6 address (format: [ipv6]:port)"
+        )
+
+    try:
+        port = int(port_str)
+    except ValueError:
+        raise ValueError(f"invalid port in IPv6 address: '{port_str}'")
+    return port, host
+
+
+def is_port_available(port):
+    """Return whether a port is available."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        try:
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            s.bind(("", port))
+            s.listen(1)
+            return True
+        except socket.error:
+            return False
+        except OverflowError:
+            return False
+
+
+def get_zmq_socket(
+    context: zmq.Context,
+    socket_type: zmq.SocketType,
+    endpoint: str,
+    bind: bool,
+    max_bind_retries: int = 10,
+    same_port: bool = False,
+) -> tuple[zmq.Socket, str]:
+    """
+    Create and configure a ZMQ socket.
+
+    Args:
+        context: ZMQ context
+        socket_type: Type of ZMQ socket
+        endpoint: Endpoint string (e.g., "tcp://localhost:5555")
+        bind: Whether to bind (True) or connect (False)
+        max_bind_retries: Maximum number of retries if bind fails due to address already in use
+        same_port: If True, retry on the same port instead of incrementing.
+            Useful when the port must be fixed (e.g., disagg sockets where
+            DiffusionServer connects to a pre-determined port).
+
+    Returns:
+        A tuple of (socket, actual_endpoint). The actual_endpoint may differ from the
+        requested endpoint if bind retry was needed (and same_port is False).
+    """
+    mem = psutil.virtual_memory()
+    total_mem = mem.total / 1024**3
+    available_mem = mem.available / 1024**3
+    if total_mem > 32 and available_mem > 16:
+        buf_size = int(0.5 * 1024**3)
+    else:
+        buf_size = -1
+
+    socket = context.socket(socket_type)
+    if endpoint.find("[") != -1:
+        socket.setsockopt(zmq.IPV6, 1)
+
+    def set_send_opt():
+        socket.setsockopt(zmq.SNDHWM, 0)
+        socket.setsockopt(zmq.SNDBUF, buf_size)
+
+    def set_recv_opt():
+        socket.setsockopt(zmq.RCVHWM, 0)
+        socket.setsockopt(zmq.RCVBUF, buf_size)
+
+    if socket_type == zmq.PUSH:
+        set_send_opt()
+    elif socket_type == zmq.PULL:
+        set_recv_opt()
+    elif socket_type in [zmq.DEALER, zmq.REQ, zmq.REP, zmq.ROUTER]:
+        set_send_opt()
+        set_recv_opt()
+    else:
+        raise ValueError(f"Unsupported socket type: {socket_type}")
+
+    if bind:
+        # Parse port from endpoint for retry logic
+        import re
+
+        port_match = re.search(r":(\d+)$", endpoint)
+
+        if port_match and max_bind_retries > 1:
+            import time as _time
+
+            original_port = int(port_match.group(1))
+            last_exception = None
+
+            for attempt in range(max_bind_retries):
+                try:
+                    current_endpoint = endpoint
+                    if attempt > 0 and not same_port:
+                        # Try next port (increment by 42 to match settle_port logic)
+                        current_port = original_port + attempt * 42
+                        current_endpoint = re.sub(
+                            r":(\d+)$", f":{current_port}", endpoint
+                        )
+                        logger.info(
+                            f"ZMQ bind failed for port {original_port + (attempt - 1) * 42}, "
+                            f"retrying with port {current_port} (attempt {attempt + 1}/{max_bind_retries})"
+                        )
+                    elif attempt > 0:
+                        logger.info(
+                            f"ZMQ bind attempt {attempt + 1}/{max_bind_retries} "
+                            f"on same port {original_port}..."
+                        )
+
+                    socket.bind(current_endpoint)
+
+                    if attempt > 0:
+                        logger.warning(
+                            f"Successfully bound ZMQ socket to {current_endpoint} after {attempt + 1} attempts. "
+                            f"Original port {original_port} was unavailable."
+                        )
+
+                    return socket, current_endpoint
+
+                except zmq.ZMQError as e:
+                    last_exception = e
+                    if e.errno == zmq.EADDRINUSE and attempt < max_bind_retries - 1:
+                        # Address already in use, retry
+                        # Longer sleep for same_port (waiting for TIME_WAIT release)
+                        _time.sleep(1.0 if same_port else 0.5)
+                        # Re-create socket since ZMQ socket state may be invalid after failed bind
+                        socket.close()
+                        socket = context.socket(socket_type)
+                        if endpoint.find("[") != -1:
+                            socket.setsockopt(zmq.IPV6, 1)
+                        if socket_type == zmq.PUSH:
+                            set_send_opt()
+                        elif socket_type == zmq.PULL:
+                            set_recv_opt()
+                        elif socket_type in [zmq.DEALER, zmq.REQ, zmq.REP, zmq.ROUTER]:
+                            set_send_opt()
+                            set_recv_opt()
+                        continue
+                    elif attempt == max_bind_retries - 1:
+                        # Last attempt failed
+                        logger.error(
+                            f"Failed to bind ZMQ socket after {max_bind_retries} attempts. "
+                            f"Original endpoint: {endpoint}, Last tried port: {original_port + attempt * 42}"
+                        )
+                        raise
+                    else:
+                        # Different error, raise immediately
+                        raise
+
+            # Should not reach here, but just in case
+            if last_exception:
+                raise last_exception
+        else:
+            # No retry logic needed (either no port in endpoint or max_bind_retries == 1)
+            socket.bind(endpoint)
+            return socket, endpoint
+    else:
+        socket.connect(endpoint)
+        return socket, endpoint
+
+    return socket, endpoint
+
+
+# https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip
+
+
+@lru_cache(maxsize=1)
+def is_host_cpu_x86() -> bool:
+    machine = platform.machine().lower()
+    return (
+        machine in ("x86_64", "amd64", "i386", "i686")
+        and hasattr(torch, "cpu")
+        and torch.cpu.is_available()
+    )
+
+
+# cuda
+
+
+def set_cuda_arch():
+    """Set CUDA architecture for compilation. Only applies to CUDA devices."""
+    if torch.cuda.is_available():
+        capability = torch.cuda.get_device_capability()
+        arch = f"{capability[0]}.{capability[1]}"
+        os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}"
+    # For XPU or other platforms, no arch setting needed
+
+
+# musa
+
+
+def set_musa_arch():
+    capability = torch.cuda.get_device_capability()
+    arch = f"{capability[0]}{capability[1]}"
+    os.environ["TORCH_MUSA_ARCH_LIST"] = f"{arch}"
+
+
+# env var managements
+
+_warned_bool_env_var_keys = set()
+
+
+def get_bool_env_var(name: str, default: str = "false") -> bool:
+    value = os.getenv(name, default)
+    value = str(value).strip().lower()
+
+    truthy_values = {"1", "true", "yes", "y", "t", "on"}
+    falsy_values = {"0", "false", "no", "n", "f", "off", ""}
+
+    if (value not in truthy_values) and (value not in falsy_values):
+        if value not in _warned_bool_env_var_keys:
+            logger.warning(
+                f"get_bool_env_var({name}) see non-understandable value={value} and treat as false"
+            )
+        _warned_bool_env_var_keys.add(value)
+
+    return value in truthy_values
+
+
+try:
+    import sgl_kernel  # noqa: F401
+
+    is_intel_amx_backend_available = hasattr(
+        torch.ops.sgl_kernel, "convert_weight_packed"
+    )
+except:
+    is_intel_amx_backend_available = False
+
+try:
+    # move torch.cpu._is_amx_tile_supported() from cpu_has_amx_support
+    # to support torch compile
+    is_amx_tile_supported = torch.cpu._is_amx_tile_supported()
+except:
+    is_amx_tile_supported = False
+
+
+def cpu_has_amx_support():
+    return is_amx_tile_supported and is_intel_amx_backend_available
+
+
+def use_intel_amx_backend(layer):
+    return getattr(layer, "use_intel_amx_backend", False)

python/sglang/multimodal_gen/runtime/utils/distributed.py

@@ -0,0 +1,234 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import pickle
+from typing import Any, List, Optional
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+
+def broadcast_pyobj(
+    data: List[Any],
+    rank: int,
+    dist_group: Optional[torch.distributed.ProcessGroup] = None,
+    src: int = 0,
+    force_cpu_device: bool = True,
+):
+    """Broadcast inputs from src rank to all other ranks with torch.dist backend.
+    The `rank` here refer to the source rank on global process group (regardless
+    of dist_group argument).
+    """
+
+    device = torch.device(
+        current_platform.device_type if not force_cpu_device else "cpu"
+    )
+
+    if rank == src:
+        if data is None or len(data) == 0:
+            tensor_size = torch.tensor([0], dtype=torch.long, device=device)
+            dist.broadcast(tensor_size, src=src, group=dist_group)
+        else:
+            serialized_data = pickle.dumps(data)
+            size = len(serialized_data)
+
+            tensor_data = torch.ByteTensor(
+                np.frombuffer(serialized_data, dtype=np.uint8).copy()
+            ).to(device)
+            tensor_size = torch.tensor([size], dtype=torch.long, device=device)
+
+            dist.broadcast(tensor_size, src=src, group=dist_group)
+            dist.broadcast(tensor_data, src=src, group=dist_group)
+        return data
+    else:
+        tensor_size = torch.tensor([0], dtype=torch.long, device=device)
+        dist.broadcast(tensor_size, src=src, group=dist_group)
+        size = tensor_size.item()
+
+        if size == 0:
+            return []
+
+        tensor_data = torch.empty(size, dtype=torch.uint8, device=device)
+        dist.broadcast(tensor_data, src=src, group=dist_group)
+
+        serialized_data = bytes(tensor_data.cpu().numpy())
+        data = pickle.loads(serialized_data)
+        return data
+
+
+def generate_masked_orthogonal_rank_groups(
+    world_size: int, parallel_size: list[int], mask: list[bool]
+) -> list[list[int]]:
+    """Generate orthogonal parallel groups based on the parallel size and mask.
+
+    Arguments:
+        world_size (int): world size
+
+        parallel_size (List[int]):
+            The parallel size of each orthogonal parallel type. For example, if
+            tensor_parallel_size = 2, pipeline_model_parallel_group = 3, data_parallel_size = 4,
+            and the parallel mapping order is tp-pp-dp, then the parallel_size = [2, 3, 4].
+
+        mask (List[bool]):
+            The mask controls which parallel methods the generated groups represent. If mask[i] is
+            True, it means the generated group contains the i-th parallelism method. For example,
+            if parallel_size = [tp_size, pp_size, dp_size], and mask = [True, False , True], then
+            the generated group is the `tp-dp` group, if the mask = [False, True, False], then the
+            generated group is the `pp` group.
+
+    Algorithm:
+        For orthogonal parallelism, such as tp/dp/pp/cp, the global_rank and
+
+        If we want to get the `dp_group` (tp_size * pp_size groups of dp_size ranks each.
+        For example,  if the gpu size is 8 and order is 'tp-pp-dp', size is '2-2-2', and the
+        dp_group here is [[0, 4], [1, 5], [2, 6], [3, 7]].)
+        The tp_rank and pp_rank will be combined to form the `dp_group_index`.
+            dp_group_index = tp_rank + pp_rank * tp_size (2)
+
+        So, Given that tp_rank and pp_rank satisfy equation (2), and dp_rank in
+        range(0, dp_size), the ranks in dp_group[dp_group_index] satisfies the
+        equation (1).
+
+        This function solve this math problem.
+
+    For example, if the parallel_size = [tp_size, dp_size, pp_size] = [2, 3, 4],
+    and the mask = [False, True, False]. Then,
+        dp_group_index(0) = tp_rank(0) + pp_rank(0) * 2
+        dp_group_index(1) = tp_rank(1) + pp_rank(0) * 2
+        ...
+        dp_group_index(7) = tp_rank(1) + pp_rank(3) * 2
+
+        dp_group[0] = 0 + range(0, 3) * 2 + 0 = [0, 2, 4]
+        dp_group[1] = 1 + range(0, 3) * 2 + 0 = [1, 3, 5]
+        ...
+        dp_group[7] = 1 + range(0, 3) * 2 + 3 * 2 * 3 = [19, 21, 23]
+    """
+
+    def prefix_product(a: List[int], init=1) -> List[int]:
+        r = [init]
+        for v in a:
+            init = init * v
+            r.append(init)
+        return r
+
+    def inner_product(a: List[int], b: List[int]) -> int:
+        return sum([x * y for x, y in zip(a, b)])
+
+    def decompose(index, shape, stride=None):
+        """
+        This function solve the math problem below:
+            There is an equation:
+                index = sum(idx[i] * stride[i])
+            And given the value of index, stride.
+            Return the idx.
+        This function will used to get the pp/dp/pp_rank
+        from group_index and rank_in_group.
+        """
+        if stride is None:
+            stride = prefix_product(shape)
+        idx = [(index // d) % s for s, d in zip(shape, stride)]
+        # stride is a prefix_product result. And the value of stride[-1]
+        # is not used.
+        assert (
+            sum([x * y for x, y in zip(idx, stride[:-1])]) == index
+        ), "idx {} with shape {} mismatch the return idx {}".format(index, shape, idx)
+        return idx
+
+    masked_shape = [s for s, m in zip(parallel_size, mask) if m]
+    unmasked_shape = [s for s, m in zip(parallel_size, mask) if not m]
+
+    global_stride = prefix_product(parallel_size)
+    masked_stride = [d for d, m in zip(global_stride, mask) if m]
+    unmasked_stride = [d for d, m in zip(global_stride, mask) if not m]
+
+    group_size = prefix_product(masked_shape)[-1]
+    num_of_group = world_size // group_size
+
+    ranks = []
+    for group_index in range(num_of_group):
+        # get indices from unmaksed for group_index.
+        decomposed_group_idx = decompose(group_index, unmasked_shape)
+        rank = []
+        for rank_in_group in range(group_size):
+            # get indices from masked for rank_in_group.
+            decomposed_rank_idx = decompose(rank_in_group, masked_shape)
+            rank.append(
+                inner_product(decomposed_rank_idx, masked_stride)
+                + inner_product(decomposed_group_idx, unmasked_stride)
+            )
+        ranks.append(rank)
+    return ranks
+
+
+class RankGenerator(object):
+    def __init__(
+        self,
+        tp: int,
+        sp: int,
+        pp: int,
+        cfg: int,
+        dp: int,
+        order: str,
+        rank_offset: int = 0,
+    ) -> None:
+        self.tp = tp
+        self.sp = sp
+        self.pp = pp
+        self.cfg = cfg
+        self.dp = dp
+        self.rank_offset = rank_offset
+        self.world_size = tp * sp * pp * cfg * dp
+
+        self.name_to_size = {
+            "tp": self.tp,
+            "sp": self.sp,
+            "pp": self.pp,
+            "cfg": self.cfg,
+            "dp": self.dp,
+        }
+        order = order.lower()
+
+        for name in self.name_to_size.keys():
+            if name not in order and self.name_to_size[name] != 1:
+                raise RuntimeError(
+                    f"The size of ({name}) is ({self.name_to_size[name]}), but you haven't specified the order ({self.order})."
+                )
+            elif name not in order:
+                order = order + "-" + name
+
+        self.order = order
+        self.ordered_size = []
+
+        for token in order.split("-"):
+            self.ordered_size.append(self.name_to_size[token])
+
+    def get_mask(self, order: str, token: str):
+        ordered_token = order.split("-")
+        token = token.split("-")
+        mask = [False] * len(ordered_token)
+        for t in token:
+            mask[ordered_token.index(t)] = True
+        return mask
+
+    def get_ranks(self, token):
+        """Get rank group by input token.
+
+        Arguments:
+            token (str):
+                Specify the ranks type that want to get. If we want
+                to obtain multiple parallel types, we can use a hyphen
+                '-' to separate them. For example, if we want to obtain
+                the TP_DP group, the token should be 'tp-dp'.
+
+        """
+        mask = self.get_mask(self.order, token)
+        ranks = generate_masked_orthogonal_rank_groups(
+            self.world_size, self.ordered_size, mask
+        )
+        if self.rank_offset > 0:
+            for rank_group in ranks:
+                for i in range(len(rank_group)):
+                    rank_group[i] += self.rank_offset
+        return ranks

python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py

@@ -0,0 +1,943 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/hf_transformers_utils.py
+
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for Huggingface Transformers."""
+
+import contextlib
+import glob
+import json
+import os
+import shutil
+import time
+from functools import reduce
+from pathlib import Path
+from typing import Any, Optional, Union, cast
+
+from diffusers.loaders.lora_base import (
+    _best_guess_weight_name,  # watch out for potetential removal from diffusers
+)
+from huggingface_hub.errors import (
+    LocalEntryNotFoundError,
+    RepositoryNotFoundError,
+    RevisionNotFoundError,
+)
+from requests.exceptions import ConnectionError as RequestsConnectionError
+from requests.exceptions import RequestException
+from transformers import AutoConfig, PretrainedConfig
+
+from sglang.multimodal_gen.runtime.loader.utils import _clean_hf_config_inplace
+from sglang.multimodal_gen.runtime.loader.weight_utils import get_lock
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.model_overlay import (
+    maybe_load_overlay_model_index,
+    maybe_resolve_overlay_model_path,
+)
+from sglang.multimodal_gen.runtime.utils.quantization_utils import (
+    normalize_flat_modelopt_quant_config,
+)
+from sglang.srt.environ import envs
+from sglang.utils import is_in_ci
+
+logger = init_logger(__name__)
+
+
+def _check_index_files_for_missing_shards(
+    model_path: str,
+) -> tuple[bool, list[str], list[str]]:
+    """
+    Check all subdirectories for missing shards based on index files.
+
+    This catches cases where a model download was interrupted, leaving
+    some safetensors shards missing while the index file exists.
+
+    Args:
+        model_path: Path to the model directory
+
+    Returns:
+        Tuple of (all_valid, missing_files, checked_subdirs)
+    """
+    missing_files = []
+    checked_subdirs = []
+
+    # Add common subdirectories for diffusers models
+    try:
+        subdirs = os.listdir(model_path)
+    except OSError as e:
+        logger.warning("Failed to list model directory %s: %s", model_path, e)
+        return True, [], []  # Assume valid if we can't check
+
+    # Check the root directory and all subdirectories that might contain model weights
+    dirs_to_check = [model_path]
+
+    for subdir in subdirs:
+        subdir_path = os.path.join(model_path, subdir)
+        if os.path.isdir(subdir_path):
+            dirs_to_check.append(subdir_path)
+
+    for dir_path in dirs_to_check:
+        # Find all safetensors index files
+        index_files = glob.glob(os.path.join(dir_path, "*.safetensors.index.json"))
+
+        for index_file in index_files:
+            checked_subdirs.append(os.path.basename(dir_path))
+            try:
+                with open(index_file) as f:
+                    index_data = json.load(f)
+
+                weight_map = index_data.get("weight_map", {})
+                if not weight_map:
+                    continue
+
+                # Get unique files referenced in weight_map
+                required_files = set(weight_map.values())
+
+                for file_name in required_files:
+                    file_path = os.path.join(dir_path, file_name)
+                    if not os.path.exists(file_path):
+                        relative_path = os.path.relpath(file_path, model_path)
+                        missing_files.append(relative_path)
+
+            except Exception as e:
+                logger.warning("Failed to read index file %s: %s", index_file, e)
+                continue
+
+    return len(missing_files) == 0, missing_files, checked_subdirs
+
+
+def _cleanup_model_cache(model_path: str, reason: str) -> bool:
+    """
+    Remove the model cache directory to force a clean re-download.
+
+    Args:
+        model_path: Path to the model directory (snapshot path)
+        reason: Reason for cleanup (for logging)
+
+    Returns:
+        True if cleanup was performed, False otherwise
+    """
+    # Navigate up to the model root directory: snapshots/hash -> snapshots -> model_root
+    # HF cache structure: models--org--name/snapshots/hash/
+    try:
+        snapshot_dir = os.path.abspath(model_path)
+        snapshots_dir = os.path.dirname(snapshot_dir)
+        repo_folder = os.path.dirname(snapshots_dir)
+
+        # Verify this looks like an HF cache structure
+        if os.path.basename(snapshots_dir) != "snapshots":
+            logger.warning(
+                "Model path %s doesn't appear to be in HF cache structure, skipping cleanup",
+                model_path,
+            )
+            return False
+
+        logger.warning(
+            "Removing model cache at %s. Reason: %s",
+            repo_folder,
+            reason,
+        )
+        shutil.rmtree(repo_folder)
+        logger.info("Successfully removed corrupted cache directory")
+        return True
+    except Exception as e:
+        logger.error(
+            "Failed to remove corrupted cache directory %s: %s. "
+            "Manual cleanup may be required.",
+            model_path,
+            e,
+        )
+        return False
+
+
+def _ci_validate_diffusers_model(model_path: str) -> tuple[bool, bool]:
+    """
+    CI-specific validation for diffusers models.
+
+    Checks all subdirectories (transformer, transformer_2, vae, etc.) for
+    missing shards based on their index files. If issues are found in CI,
+    cleans up the cache to force re-download.
+
+    Args:
+        model_path: Path to the model directory
+
+    Returns:
+        Tuple of (is_valid, cleanup_performed)
+        - is_valid: True if the model is valid
+        - cleanup_performed: True if cleanup was performed (only relevant when is_valid=False)
+    """
+    if not is_in_ci():
+        return True, False
+    is_valid, missing_files, checked_subdirs = _check_index_files_for_missing_shards(
+        model_path
+    )
+
+    if not is_valid:
+        logger.error(
+            "CI validation failed for %s. Missing %d file(s): %s. "
+            "Checked subdirectories: %s",
+            model_path,
+            len(missing_files),
+            missing_files[:5] if len(missing_files) > 5 else missing_files,
+            checked_subdirs,
+        )
+        cleanup_performed = _cleanup_model_cache(
+            model_path,
+            f"Missing {len(missing_files)} shard file(s): {missing_files[:3]}",
+        )
+        return False, cleanup_performed
+
+    if checked_subdirs:
+        logger.info(
+            "CI validation passed for %s. Checked subdirectories: %s",
+            model_path,
+            checked_subdirs,
+        )
+
+    return True, False
+
+
+def _verify_diffusers_model_complete(path: str) -> bool:
+    """Check if a diffusers model directory has all required component subdirectories."""
+    config_path = os.path.join(path, "model_index.json")
+    if not os.path.exists(config_path):
+        return False
+
+    try:
+        with open(config_path) as config_file:
+            model_index = json.load(config_file)
+    except Exception as exc:
+        logger.warning("Failed to read model_index.json at %s: %s", config_path, exc)
+        return False
+
+    component_keys = [
+        key
+        for key, value in model_index.items()
+        if isinstance(value, (list, tuple))
+        and len(value) == 2
+        and all(isinstance(item, str) for item in value)
+    ]
+    if component_keys:
+        return all(os.path.exists(os.path.join(path, key)) for key in component_keys)
+
+    return os.path.exists(os.path.join(path, "transformer")) and os.path.exists(
+        os.path.join(path, "vae")
+    )
+
+
+_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = {
+    # ChatGLMConfig.model_type: ChatGLMConfig,
+    # DbrxConfig.model_type: DbrxConfig,
+    # ExaoneConfig.model_type: ExaoneConfig,
+    # Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig,
+}
+
+for name, cls in _CONFIG_REGISTRY.items():
+    with contextlib.suppress(ValueError):
+        AutoConfig.register(name, cls)
+
+
+def download_from_hf(model_path: str):
+    if os.path.exists(model_path):
+        return model_path
+
+    return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"])
+
+
+def get_hf_config(
+    component_model_path: str,
+    trust_remote_code: bool,
+    revision: str | None = None,
+    model_override_args: dict | None = None,
+    **kwargs,
+) -> PretrainedConfig:
+    if check_gguf_file(component_model_path):
+        raise NotImplementedError("GGUF models are not supported.")
+
+    config = AutoConfig.from_pretrained(
+        component_model_path,
+        trust_remote_code=trust_remote_code,
+        revision=revision,
+        **kwargs,
+    )
+    if config.model_type in _CONFIG_REGISTRY:
+        config_class = _CONFIG_REGISTRY[config.model_type]
+        config = config_class.from_pretrained(component_model_path, revision=revision)
+        # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`.
+        config._name_or_path = component_model_path
+    if model_override_args:
+        config.update(model_override_args)
+
+    return config
+
+
+def get_config(
+    model: str,
+    trust_remote_code: bool,
+    revision: Optional[str] = None,
+    model_override_args: Optional[dict] = None,
+    **kwargs,
+):
+    return AutoConfig.from_pretrained(
+        model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
+    )
+
+
+def load_dict(file_path):
+    if not os.path.exists(file_path):
+        return {}
+    try:
+        # Load the config directly from the file
+        with open(file_path) as f:
+            config_dict: dict[str, Any] = json.load(f)
+        if "_diffusers_version" in config_dict:
+            config_dict.pop("_diffusers_version")
+        # TODO(will): apply any overrides from inference args
+        return config_dict
+    except Exception as e:
+        raise RuntimeError(
+            f"Failed to load diffusers config from {file_path}: {e}"
+        ) from e
+
+
+def prepare_diffusers_component_path_for_loading(component_path: str) -> str:
+    """Download component repos if needed and patch legacy flat ModelOpt configs."""
+    local_component_path = (
+        maybe_download_model(component_path)
+        if not os.path.exists(component_path)
+        else component_path
+    )
+    config_path = os.path.join(local_component_path, "config.json")
+    if not os.path.exists(config_path):
+        return local_component_path
+
+    with get_lock(config_path):
+        try:
+            with open(config_path, encoding="utf-8") as f:
+                config = cast(dict[str, Any], json.load(f))
+        except Exception as exc:
+            logger.warning("Failed to read component config %s: %s", config_path, exc)
+            return local_component_path
+
+        quant_config = config.get("quantization_config")
+        normalized_quant_config = normalize_flat_modelopt_quant_config(quant_config)
+        if normalized_quant_config == quant_config:
+            return local_component_path
+
+        config["quantization_config"] = normalized_quant_config
+        try:
+            with open(config_path, "w", encoding="utf-8") as f:
+                json.dump(config, f, indent=2, sort_keys=True)
+                f.write("\n")
+        except OSError as exc:
+            logger.warning(
+                "Could not persist normalized ModelOpt config at %s (%s); "
+                "normalization will be applied in memory at load time.",
+                config_path,
+                exc,
+            )
+        else:
+            logger.warning(
+                "Patched legacy flat ModelOpt quantization_config at %s with quant_type=%s "
+                "for diffusers compatibility.",
+                config_path,
+                normalized_quant_config.get("quant_type"),
+            )
+
+    return local_component_path
+
+
+def get_diffusers_component_config(
+    component_path: str,
+) -> dict[str, Any]:
+    """Gets a configuration of a submodule for the given diffusers model."""
+    # Download from HuggingFace Hub if path doesn't exist locally
+    component_path = prepare_diffusers_component_path_for_loading(component_path)
+
+    config_names = ["generation_config.json"]
+    # By default, we load config.json, but scheduler_config.json for scheduler
+    if "scheduler" in component_path:
+        config_names.append("scheduler_config.json")
+    else:
+        config_names.append("config.json")
+
+    config_file_paths = [
+        os.path.join(component_path, config_name) for config_name in config_names
+    ]
+
+    combined_config = reduce(
+        lambda acc, path: acc | load_dict(path), config_file_paths, {}
+    )
+
+    quant_config = combined_config.get("quantization_config")
+    if quant_config is not None:
+        combined_config["quantization_config"] = normalize_flat_modelopt_quant_config(
+            quant_config
+        )
+
+    _clean_hf_config_inplace(combined_config)
+
+    logger.debug("HF model config: %s", combined_config)
+
+    return combined_config
+
+
+# Models don't use the same configuration key for determining the maximum
+# context length.  Store them here so we can sanely check them.
+# NOTE: The ordering here is important. Some models have two of these and we
+# have a preference for which value gets used.
+CONTEXT_LENGTH_KEYS = [
+    "max_sequence_length",
+    "seq_length",
+    "max_seq_len",
+    "model_max_length",
+    "max_position_embeddings",
+]
+
+
+def attach_additional_stop_token_ids(tokenizer):
+    # Special handling for stop token <|eom_id|> generated by llama 3 tool use.
+    if "<|eom_id|>" in tokenizer.get_added_vocab():
+        tokenizer.additional_stop_token_ids = {
+            tokenizer.get_added_vocab()["<|eom_id|>"]
+        }
+    else:
+        tokenizer.additional_stop_token_ids = None
+
+
+def check_gguf_file(model: str | os.PathLike) -> bool:
+    """Check if the file is a GGUF model."""
+    model = Path(model)
+    if not model.is_file():
+        return False
+    elif model.suffix == ".gguf":
+        return True
+
+    with open(model, "rb") as f:
+        header = f.read(4)
+    return header == b"GGUF"
+
+
+def maybe_download_lora(
+    model_name_or_path: str,
+    local_dir: str | None = None,
+    download: bool = True,
+    weight_name: str | None = None,
+) -> str:
+    """
+    Check if the model path is a Hugging Face Hub model ID and download it if needed.
+    Args:
+        model_name_or_path: Local path or Hugging Face Hub model ID
+        local_dir: Local directory to save the model
+        download: Whether to download the model from Hugging Face Hub
+        weight_name: Specific safetensors filename to load (pins deterministic selection
+                     for repos with multiple weight files)
+
+    Returns:
+        Local path to the model
+    """
+    allow_patterns = ["*.json", "*.safetensors", "*.bin"]
+
+    local_path = maybe_download_model(
+        model_name_or_path,
+        local_dir,
+        download,
+        is_lora=True,
+        allow_patterns=allow_patterns,
+    )
+    # return directly if local_path is a file
+    if os.path.isfile(local_path):
+        return local_path
+
+    if weight_name is not None:
+        target = os.path.join(local_path, weight_name)
+        if not os.path.isfile(target):
+            raise FileNotFoundError(
+                f"Specified lora_weight_name '{weight_name}' not found in {local_path}"
+            )
+        return target
+
+    guessed = _best_guess_weight_name(local_path, file_extension=".safetensors")
+    # AMD workaround: PR 15813 changed from model_name_or_path to local_path,
+    # which can return None. Fall back to original behavior on ROCm.
+    if guessed is None and current_platform.is_rocm():
+        guessed = _best_guess_weight_name(
+            model_name_or_path, file_extension=".safetensors"
+        )
+    return os.path.join(local_path, guessed)
+
+
+def verify_model_config_and_directory(model_path: str) -> dict[str, Any]:
+    """
+    Verify that the model directory contains a valid diffusers configuration.
+
+    Args:
+        model_path: Path to the model directory
+
+    Returns:
+        The loaded model configuration as a dictionary
+    """
+
+    # Check for model_index.json which is required for diffusers models
+    config_path = os.path.join(model_path, "model_index.json")
+    if not os.path.exists(config_path):
+        raise ValueError(
+            f"Model directory {model_path} does not contain model_index.json. "
+            "Only HuggingFace diffusers format is supported."
+        )
+
+    # Load the config
+    with open(config_path) as f:
+        config = json.load(f)
+
+    # Verify diffusers version exists
+    if "_diffusers_version" not in config:
+        raise ValueError("model_index.json does not contain _diffusers_version")
+
+    logger.info("Diffusers version: %s", config["_diffusers_version"])
+
+    component_keys = [
+        key
+        for key, value in config.items()
+        if isinstance(value, (list, tuple))
+        and len(value) == 2
+        and all(isinstance(item, str) for item in value)
+    ]
+    if component_keys:
+        missing_components = [
+            component_key
+            for component_key in component_keys
+            if not os.path.exists(os.path.join(model_path, component_key))
+        ]
+        if missing_components:
+            missing_str = ", ".join(missing_components)
+            raise ValueError(
+                f"Model directory {model_path} is missing required component "
+                f"directories: {missing_str}."
+            )
+    else:
+        transformer_dir = os.path.join(model_path, "transformer")
+        vae_dir = os.path.join(model_path, "vae")
+        if not os.path.exists(transformer_dir):
+            raise ValueError(
+                f"Model directory {model_path} does not contain a transformer/ directory."
+            )
+        if not os.path.exists(vae_dir):
+            raise ValueError(
+                f"Model directory {model_path} does not contain a vae/ directory."
+            )
+    return cast(dict[str, Any], config)
+
+
+def _resolve_remote_repo_model_index_path(model_name_or_path: str) -> str:
+    """Return a local path to a remote repo's ``model_index.json``"""
+    from huggingface_hub.errors import EntryNotFoundError
+
+    try:
+        # Cache-aware: no local_dir, so HF reuses the cache and revalidates the
+        # ETag against the Hub, re-downloading only when the remote changed.
+        return hf_hub_download(repo_id=model_name_or_path, filename="model_index.json")
+    except EntryNotFoundError:
+        # Repo exists but has no model_index.json (single-model repo); let the
+        # caller fall through to the single-model path.
+        raise
+    except Exception as online_err:
+        cached_path = None
+        if not envs.SGLANG_USE_MODELSCOPE.get():
+            from huggingface_hub import try_to_load_from_cache
+
+            cached = try_to_load_from_cache(
+                repo_id=model_name_or_path, filename="model_index.json"
+            )
+            if isinstance(cached, str) and os.path.exists(cached):
+                cached_path = cached
+        if cached_path is not None:
+            logger.warning(
+                "Could not fetch model_index.json for '%s' from the Hugging Face "
+                "Hub (%s); using the locally cached copy at '%s'. The cached copy "
+                "may be out of date — provide an HF token or clear the cache to "
+                "force a refresh.",
+                model_name_or_path,
+                online_err,
+                cached_path,
+            )
+            return cached_path
+        raise
+
+
+def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]:
+    """
+    Download and extract just the model_index.json for a Hugging Face model.
+
+    Args:
+        model_name_or_path: Path or HF Hub model ID
+
+    Returns:
+        The parsed model_index.json as a dictionary
+    """
+    from huggingface_hub.errors import EntryNotFoundError
+
+    overlay_config = maybe_load_overlay_model_index(
+        model_name_or_path,
+        snapshot_download_fn=snapshot_download,
+        hf_hub_download_fn=hf_hub_download,
+    )
+    if overlay_config is not None:
+        return overlay_config
+
+    # If it's a local path, verify it directly.
+    if os.path.exists(model_name_or_path):
+        try:
+            return verify_model_config_and_directory(model_name_or_path)
+        except ValueError:
+            # Not a pipeline, maybe a single model.
+            config_path = os.path.join(model_name_or_path, "config.json")
+            if os.path.exists(config_path):
+                with open(config_path) as f:
+                    config = json.load(f)
+                return config
+            raise
+
+    # For remote models, resolve model_index.json (Hub-first, cache fallback).
+    try:
+        model_index_path = _resolve_remote_repo_model_index_path(model_name_or_path)
+
+        # Load the model_index.json
+        with open(model_index_path) as f:
+            config: dict[str, Any] = json.load(f)
+
+        # Verify it has the required fields
+        if "_class_name" not in config:
+            raise ValueError(
+                f"model_index.json for {model_name_or_path} does not contain _class_name field"
+            )
+
+        if "_diffusers_version" not in config:
+            raise ValueError(
+                f"model_index.json for {model_name_or_path} does not contain _diffusers_version field"
+            )
+
+        # Add the pipeline name for downstream use
+        config["pipeline_name"] = config["_class_name"]
+
+        logger.debug(
+            "Resolved model_index.json for %s, pipeline: %s",
+            model_name_or_path,
+            config["_class_name"],
+        )
+        return config
+    except EntryNotFoundError:
+        logger.debug(
+            "model_index.json not found for %s. Assuming it is a single model and downloading it.",
+            model_name_or_path,
+        )
+        local_path = maybe_download_model(model_name_or_path)
+        config_path = os.path.join(local_path, "config.json")
+        if not os.path.exists(config_path):
+            raise ValueError(
+                f"Failed to find config.json for {model_name_or_path} after failing to find model_index.json"
+                f"You might be looking for models ending with '-Diffusers'"
+            )
+        with open(config_path) as f:
+            config = json.load(f)
+        return config
+    except Exception as e:
+        raise ValueError(
+            f"Failed to download or parse model_index.json for {model_name_or_path}: {e}"
+        ) from e
+
+
+def maybe_download_model(
+    model_name_or_path: str,
+    local_dir: str | None = None,
+    download: bool = True,
+    is_lora: bool = False,
+    allow_patterns: list[str] | None = None,
+    force_diffusers_model: bool = False,
+    skip_overlay_resolution: bool = False,
+) -> str:
+    """
+    Check if the model path is a Hugging Face Hub model ID and download it if needed.
+
+    Args:
+        model_name_or_path: Local path or Hugging Face Hub model ID
+        local_dir: Local directory to save the model
+        download: Whether to download the model from Hugging Face Hub
+        is_lora: If True, skip model completeness verification (LoRA models don't have transformer/vae directories)
+        force_diffusers_model: If True, apply diffusers model check. Otherwise it should be a component model
+    Returns:
+        Local path to the model
+    """
+    if force_diffusers_model and not skip_overlay_resolution:
+        # return overlay model path if applicable
+        overlay_model_path = maybe_resolve_overlay_model_path(
+            model_name_or_path,
+            local_dir=local_dir,
+            download=download,
+            allow_patterns=allow_patterns,
+            snapshot_download_fn=snapshot_download,
+            hf_hub_download_fn=hf_hub_download,
+            verify_diffusers_model_complete_fn=_verify_diffusers_model_complete,
+            base_model_download_fn=maybe_download_model,
+        )
+        if overlay_model_path is not None:
+            return overlay_model_path
+
+    # 1. Local path check: if path exists locally, verify it's complete (skip for LoRA)
+    if os.path.exists(model_name_or_path):
+        if not force_diffusers_model:
+            return model_name_or_path
+        if is_lora or _verify_diffusers_model_complete(model_name_or_path):
+            if not is_lora:
+                is_valid, cleanup_performed = _ci_validate_diffusers_model(
+                    model_name_or_path
+                )
+                if not is_valid:
+                    if cleanup_performed:
+                        logger.warning(
+                            "CI validation failed for local model at %s, "
+                            "cache has been cleaned up, will re-download",
+                            model_name_or_path,
+                        )
+                        # Fall through to download
+                    else:
+                        raise ValueError(
+                            f"CI validation failed for local model at {model_name_or_path}. "
+                            "Some safetensors shards are missing. "
+                            "Please manually delete the model directory and retry."
+                        )
+                else:
+                    logger.info("Model already exists locally and is complete")
+                    return model_name_or_path
+            else:
+                logger.info("Model already exists locally and is complete")
+                return model_name_or_path
+        else:
+            logger.warning(
+                "Local model at %s appears incomplete (missing required components), "
+                "will attempt re-download",
+                model_name_or_path,
+            )
+
+    # 2. Cache-first strategy (Fast Path)
+    # Try to read from HF cache without network access
+    try:
+        logger.info(
+            "Checking for cached model in HF Hub cache for %s...", model_name_or_path
+        )
+        local_path = snapshot_download(
+            repo_id=model_name_or_path,
+            ignore_patterns=["*.onnx", "*.msgpack"],
+            local_dir=local_dir,
+            local_files_only=True,
+            max_workers=8,
+        )
+        if not force_diffusers_model:
+            return str(local_path)
+        if is_lora or _verify_diffusers_model_complete(local_path):
+            if not is_lora:
+                is_valid, cleanup_performed = _ci_validate_diffusers_model(local_path)
+                if not is_valid:
+                    logger.warning(
+                        "CI validation failed for cached model at %s, "
+                        "%s, will re-download",
+                        local_path,
+                        (
+                            "cache has been cleaned up"
+                            if cleanup_performed
+                            else "cleanup was not performed"
+                        ),
+                    )
+                    # Fall through to download
+                else:
+                    logger.info("Found complete model in cache at %s", local_path)
+                    return str(local_path)
+            else:
+                logger.info("Found complete model in cache at %s", local_path)
+                return str(local_path)
+        else:
+            if not download:
+                raise ValueError(
+                    f"Model {model_name_or_path} found in cache but is incomplete and download=False."
+                )
+            logger.info(
+                "Model found in cache but incomplete, will download from HF Hub"
+            )
+    except LocalEntryNotFoundError:
+        if not download:
+            raise ValueError(
+                f"Model {model_name_or_path} not found in local cache and download=False."
+            )
+        logger.info("Model not found in cache, will download from HF Hub")
+    except Exception as e:
+        logger.warning(
+            "Unexpected error while checking cache for %s: %s, will attempt download",
+            model_name_or_path,
+            e,
+        )
+        if not download:
+            raise ValueError(
+                f"Error checking cache for {model_name_or_path} and download=False: {e}"
+            ) from e
+
+    # 3. Download strategy (with retry mechanism)
+    MAX_RETRIES = 5
+    for attempt in range(MAX_RETRIES):
+        try:
+            logger.info(
+                "Downloading model snapshot from HF Hub for %s (attempt %d/%d)...",
+                model_name_or_path,
+                attempt + 1,
+                MAX_RETRIES,
+            )
+            with get_lock(model_name_or_path).acquire(poll_interval=2):
+                local_path = snapshot_download(
+                    repo_id=model_name_or_path,
+                    ignore_patterns=["*.onnx", "*.msgpack"],
+                    allow_patterns=allow_patterns,
+                    local_dir=local_dir,
+                    max_workers=8,
+                )
+
+            if not force_diffusers_model:
+                return str(local_path)
+            # Verify downloaded model is complete (skip for LoRA)
+            elif not is_lora and not _verify_diffusers_model_complete(local_path):
+                logger.warning(
+                    "Downloaded model at %s is incomplete, retrying with force_download=True",
+                    local_path,
+                )
+                with get_lock(model_name_or_path).acquire(poll_interval=2):
+                    local_path = snapshot_download(
+                        repo_id=model_name_or_path,
+                        ignore_patterns=["*.onnx", "*.msgpack"],
+                        local_dir=local_dir,
+                        max_workers=8,
+                        force_download=True,
+                    )
+                if not _verify_diffusers_model_complete(local_path):
+                    raise ValueError(
+                        f"Downloaded model at {local_path} is still incomplete after forced re-download. "
+                        "The model repository may be missing required components (model_index.json, transformer/, or vae/)."
+                    )
+
+            # CI validation: check all subdirectories for missing shards after download
+            if not is_lora:
+                is_valid, cleanup_performed = _ci_validate_diffusers_model(local_path)
+                if not is_valid:
+                    # In CI, if validation fails after download, we have a serious issue
+                    # If cleanup was performed, the next retry should get a fresh download
+                    raise ValueError(
+                        f"CI validation failed for downloaded model at {local_path}. "
+                        f"Some safetensors shards are missing. Cleanup performed: {cleanup_performed}."
+                    )
+
+            logger.info("Downloaded model to %s", local_path)
+            return str(local_path)
+
+        except (RepositoryNotFoundError, RevisionNotFoundError) as e:
+            raise ValueError(
+                f"Model or revision not found at {model_name_or_path}. "
+                f"Please check the model ID or ensure you have access to the repository. Error: {e}"
+            ) from e
+        except (RequestException, RequestsConnectionError) as e:
+            if attempt == MAX_RETRIES - 1:
+                raise ValueError(
+                    f"Could not find model at {model_name_or_path} and failed to download from HF Hub "
+                    f"after {MAX_RETRIES} attempts due to network error: {e}"
+                ) from e
+            wait_time = 2**attempt
+            logger.warning(
+                "Download failed (attempt %d/%d) due to network error: %s. "
+                "Retrying in %d seconds...",
+                attempt + 1,
+                MAX_RETRIES,
+                e,
+                wait_time,
+            )
+            time.sleep(wait_time)
+        except Exception as e:
+            raise ValueError(
+                f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}"
+            ) from e
+
+
+# Unified download functions with Hugging Face-compatible names
+def hf_hub_download(
+    repo_id: str,
+    filename: str,
+    local_dir: Optional[Union[str, Path]] = None,
+    **kwargs,
+) -> str:
+    """Unified hf_hub_download that supports both Hugging Face Hub and ModelScope."""
+    if envs.SGLANG_USE_MODELSCOPE.get():
+        from modelscope import model_file_download
+
+        return model_file_download(
+            model_id=repo_id,
+            file_path=filename,
+            cache_dir=local_dir,
+            **kwargs,
+        )
+    else:
+        from huggingface_hub import hf_hub_download as _hf_hub_download
+
+        return _hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            local_dir=local_dir,
+            **kwargs,
+        )
+
+
+def snapshot_download(
+    repo_id: str,
+    local_dir: Optional[Union[str, Path]] = None,
+    ignore_patterns: Optional[Union[list[str], str]] = None,
+    allow_patterns: Optional[Union[list[str], str]] = None,
+    local_files_only: bool = False,
+    max_workers: int = 8,
+    **kwargs,
+) -> str:
+    """Unified snapshot_download that supports both Hugging Face Hub and ModelScope."""
+    if envs.SGLANG_USE_MODELSCOPE.get():
+        from modelscope import snapshot_download as _ms_snapshot_download
+
+        ms_kwargs = {
+            "model_id": repo_id,
+            "local_dir": local_dir,
+            "ignore_patterns": ignore_patterns,
+            "allow_patterns": allow_patterns,
+            "local_files_only": local_files_only,
+            "max_workers": max_workers,
+        }
+        ms_kwargs.update(kwargs)
+        return _ms_snapshot_download(**ms_kwargs)
+    else:
+        from huggingface_hub import snapshot_download as _hf_snapshot_download
+
+        hf_kwargs = {
+            "repo_id": repo_id,
+            "local_dir": local_dir,
+            "ignore_patterns": ignore_patterns,
+            "allow_patterns": allow_patterns,
+            "local_files_only": local_files_only,
+            "max_workers": max_workers,
+            "etag_timeout": 60,
+        }
+        hf_kwargs.update(kwargs)
+        return _hf_snapshot_download(**hf_kwargs)

python/sglang/multimodal_gen/runtime/utils/image_io.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+import base64
+import os
+import re
+
+
+def save_base64_image_to_path(base64_data: str, target_path: str) -> str:
+    b64_format_hint = (
+        "Failed to decode base64 image. "
+        "Expected format: `data:[<media-type>];base64,<data>`"
+    )
+
+    match = re.match(r"data:(.*?)(;base64)?,(.*)", base64_data)
+    if not match:
+        raise ValueError(b64_format_hint)
+    media_type = match.group(1)
+    is_base64 = match.group(2)
+    if not is_base64:
+        raise ValueError(f"{b64_format_hint} (missing ;base64 marker)")
+    data = match.group(3)
+    if not data:
+        raise ValueError(f"{b64_format_hint} (empty data payload)")
+
+    if media_type.startswith("image/"):
+        ext = media_type.split("/")[-1].lower()
+        if ext == "jpeg":
+            ext = "jpg"
+    else:
+        ext = "jpg"
+    target_path = f"{target_path}.{ext}"
+    os.makedirs(os.path.dirname(target_path), exist_ok=True)
+
+    try:
+        image_data = base64.b64decode(data)
+    except Exception as exc:
+        raise Exception(f"Failed to decode base64 image: {str(exc)}") from exc
+
+    with open(target_path, "wb") as f:
+        f.write(image_data)
+
+    return target_path

python/sglang/multimodal_gen/runtime/utils/logging_utils.py

@@ -0,0 +1,658 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logger.py
+"""Logging configuration for sglang.multimodal_gen."""
+
+import argparse
+import contextlib
+import dataclasses
+import datetime
+import inspect
+import logging
+import os
+import sys
+import time
+from contextlib import contextmanager
+from enum import Enum
+from functools import lru_cache, partial
+from logging import Logger
+from types import MethodType
+from typing import Any, cast
+
+import sglang.multimodal_gen.envs as envs
+
+SGLANG_DIFFUSION_LOGGING_LEVEL = envs.SGLANG_DIFFUSION_LOGGING_LEVEL
+SGLANG_DIFFUSION_LOGGING_PREFIX = envs.SGLANG_DIFFUSION_LOGGING_PREFIX
+
+# color
+CYAN = "\033[1;36m"
+RED = "\033[91m"
+GREEN = "\033[92m"
+YELLOW = "\033[93m"
+RESET = "\033[0;0m"
+
+_FORMAT = (
+    f"{SGLANG_DIFFUSION_LOGGING_PREFIX}%(levelname)s %(asctime)s "
+    "[%(filename)s: %(lineno)d] %(message)s"
+)
+
+# _FORMAT = "[%(asctime)s] %(message)s"
+_DATE_FORMAT = "%m-%d %H:%M:%S"
+
+DEFAULT_LOGGING_CONFIG = {
+    "formatters": {
+        "sgl_diffusion": {
+            "class": "sglang.multimodal_gen.runtime.utils.logging_utils.ColoredFormatter",
+            "datefmt": _DATE_FORMAT,
+            "format": _FORMAT,
+        },
+    },
+    "handlers": {
+        "sgl_diffusion": {
+            "class": "logging.StreamHandler",
+            "formatter": "sgl_diffusion",
+            "level": SGLANG_DIFFUSION_LOGGING_LEVEL,
+            "stream": "ext://sys.stdout",
+        },
+    },
+    "loggers": {
+        "sgl_diffusion": {
+            "handlers": ["sgl_diffusion"],
+            "level": "WARNING",
+            "propagate": False,
+        },
+    },
+    "root": {
+        "handlers": ["sgl_diffusion"],
+        "level": "DEBUG",
+    },
+    "version": 1,
+    "disable_existing_loggers": False,
+}
+
+
+class ColoredFormatter(logging.Formatter):
+    """A logging formatter that adds color to log levels."""
+
+    LEVEL_COLORS = {
+        logging.ERROR: RED,
+        logging.WARNING: YELLOW,
+    }
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Adds color to the log"""
+
+        formatted_message = super().format(record)
+
+        color = self.LEVEL_COLORS.get(record.levelno)
+        if color:
+            formatted_message = f"{color}{formatted_message}{RESET}"
+
+        return formatted_message
+
+
+class SortedHelpFormatter(argparse.HelpFormatter):
+    """SortedHelpFormatter that sorts arguments by their option strings."""
+
+    def add_arguments(self, actions):
+        actions = sorted(actions, key=lambda x: x.option_strings)
+        super().add_arguments(actions)
+
+
+@lru_cache
+def _print_info_once(logger: Logger, msg: str) -> None:
+    # Set the stacklevel to 2 to print the original caller's line info
+    logger.info(msg, stacklevel=2)
+
+
+@lru_cache
+def _print_warning_once(logger: Logger, msg: str) -> None:
+    # Set the stacklevel to 2 to print the original caller's line info
+    logger.warning(msg, stacklevel=2)
+
+
+def get_is_main_process():
+    try:
+        rank = int(os.environ["RANK"])
+    except (KeyError, ValueError):
+        rank = 0
+    return rank == 0
+
+
+def get_is_local_main_process():
+    try:
+        rank = int(os.environ["LOCAL_RANK"])
+    except (KeyError, ValueError):
+        rank = 0
+    return rank == 0
+
+
+def _log_process_aware(
+    server_log_level: int,
+    level: int,
+    logger_self: Logger,
+    msg: object,
+    *args: Any,
+    main_process_only: bool,
+    local_main_process_only: bool,
+    **kwargs: Any,
+) -> None:
+    """Helper function to log a message if the process rank matches the criteria."""
+    is_main_process = get_is_main_process()
+    is_local_main_process = get_is_local_main_process()
+    should_log = (
+        not main_process_only
+        and not local_main_process_only
+        or (main_process_only and is_main_process)
+        or (local_main_process_only and is_local_main_process)
+        or server_log_level <= logging.DEBUG
+    )
+
+    if should_log:
+        # stacklevel=3 to show the original caller's location,
+        # as this function is called by the patched methods.
+        if "stacklevel" in kwargs:
+            logger_self.log(level, msg, *args, **kwargs)
+        else:
+            logger_self.log(level, msg, *args, stacklevel=3, **kwargs)
+
+
+class _SGLDiffusionLogger(Logger):
+    """
+    Note:
+        This class is just to provide type information.
+        We actually patch the methods directly on the :class:`logging.Logger`
+        instance to avoid conflicting with other libraries such as
+        `intel_extension_for_pytorch.utils._logger`.
+    """
+
+    def info_once(self, msg: str) -> None:
+        """
+        As :meth:`info`, but subsequent calls with the same message
+        are silently dropped.
+        """
+        _print_info_once(self, msg)
+
+    def warning_once(self, msg: str) -> None:
+        """
+        As :meth:`warning`, but subsequent calls with the same message
+        are silently dropped.
+        """
+        _print_warning_once(self, msg)
+
+    def info(  # type: ignore[override]
+        self,
+        msg: object,
+        *args: Any,
+        main_process_only: bool = True,
+        local_main_process_only: bool = True,
+        **kwargs: Any,
+    ) -> None: ...
+
+    def debug(  # type: ignore[override]
+        self,
+        msg: object,
+        *args: Any,
+        main_process_only: bool = True,
+        local_main_process_only: bool = True,
+        **kwargs: Any,
+    ) -> None: ...
+
+    def warning(  # type: ignore[override]
+        self,
+        msg: object,
+        *args: Any,
+        main_process_only: bool = False,
+        local_main_process_only: bool = True,
+        **kwargs: Any,
+    ) -> None: ...
+
+    def error(  # type: ignore[override]
+        self,
+        msg: object,
+        *args: Any,
+        main_process_only: bool = False,
+        local_main_process_only: bool = True,
+        **kwargs: Any,
+    ) -> None: ...
+
+
+def init_logger(name: str) -> _SGLDiffusionLogger:
+    """The main purpose of this function is to ensure that loggers are
+    retrieved in such a way that we can be sure the root sgl_diffusion logger has
+    already been configured."""
+
+    logger = logging.getLogger(name)
+
+    server_log_level = logger.getEffectiveLevel()
+
+    # Patch instance methods
+    setattr(logger, "info_once", MethodType(_print_info_once, logger))
+    setattr(logger, "warning_once", MethodType(_print_warning_once, logger))
+
+    def _create_patched_method(
+        level: int,
+        main_process_only_default: bool,
+        local_main_process_only_default: bool,
+    ):
+        def _method(
+            self: Logger,
+            msg: object,
+            *args: Any,
+            main_process_only: bool = main_process_only_default,
+            local_main_process_only: bool = local_main_process_only_default,
+            **kwargs: Any,
+        ) -> None:
+            _log_process_aware(
+                server_log_level,
+                level,
+                self,
+                msg,
+                *args,
+                main_process_only=main_process_only,
+                local_main_process_only=local_main_process_only,
+                **kwargs,
+            )
+
+        return _method
+
+    setattr(
+        logger,
+        "info",
+        MethodType(_create_patched_method(logging.INFO, True, True), logger),
+    )
+    setattr(
+        logger,
+        "debug",
+        MethodType(_create_patched_method(logging.DEBUG, True, True), logger),
+    )
+    setattr(
+        logger,
+        "warning",
+        MethodType(_create_patched_method(logging.WARNING, False, True), logger),
+    )
+    setattr(
+        logger,
+        "error",
+        MethodType(_create_patched_method(logging.ERROR, False, False), logger),
+    )
+
+    return cast(_SGLDiffusionLogger, logger)
+
+
+logger = init_logger(__name__)
+
+
+def _is_torch_tensor(obj: Any) -> tuple[bool, Any]:
+    """Return (is_tensor, torch_module_or_None) without importing torch at module import time."""
+    try:
+        import torch  # type: ignore
+
+        return isinstance(obj, torch.Tensor), torch
+    except Exception:
+        return False, None
+
+
+def _sanitize_for_logging(obj: Any, key_hint: str | None = None) -> Any:
+    """Recursively convert objects to JSON-serializable forms for concise logging.
+
+    Rules:
+    - Drop any field/dict key named 'param_names_mapping'.
+    - Render Enums using their value.
+    - Render torch.Tensor as a compact summary; if key name is 'scaling_factor', include stats.
+    - Dataclasses are expanded to dicts and sanitized recursively.
+    - Callables/functions are rendered as their qualified name.
+    - Redact sensitive fields like 'prompt' and 'negative_prompt' (only show length).
+    - Fallback to str(...) for unknown types.
+    """
+    if obj is None or isinstance(obj, (str, int, float, bool)):
+        if key_hint in ("prompt", "negative_prompt"):
+            if isinstance(obj, str):
+                return f"<redacted, len={len(obj)}>"
+        return obj
+
+    if isinstance(obj, Enum):
+        return obj.value
+
+    is_tensor, torch_mod = _is_torch_tensor(obj)
+    if is_tensor:
+        try:
+            ten = obj.detach().cpu()
+            if key_hint == "scaling_factor":
+                stats = {
+                    "shape": list(ten.shape),
+                    "dtype": str(ten.dtype),
+                }
+                try:
+                    stats["min"] = float(ten.min().item())
+                except Exception:
+                    pass
+                try:
+                    stats["max"] = float(ten.max().item())
+                except Exception:
+                    pass
+                try:
+                    stats["mean"] = float(ten.float().mean().item())
+                except Exception:
+                    pass
+                return {"tensor": "scaling_factor", **stats}
+            return {"tensor": True, "shape": list(ten.shape), "dtype": str(ten.dtype)}
+        except Exception:
+            return "<tensor>"
+
+    if dataclasses.is_dataclass(obj):
+        result: dict[str, Any] = {}
+        for f in dataclasses.fields(obj):
+            if not f.repr:
+                continue
+            name = f.name
+            if "names_mapping" in name:
+                continue
+            try:
+                value = getattr(obj, name)
+            except Exception:
+                continue
+            result[name] = _sanitize_for_logging(value, key_hint=name)
+        return result
+
+    if isinstance(obj, dict):
+        result_dict: dict[str, Any] = {}
+        for k, v in obj.items():
+            try:
+                key_str = str(k)
+            except Exception:
+                key_str = "<key>"
+            if key_str == "param_names_mapping":
+                continue
+            result_dict[key_str] = _sanitize_for_logging(v, key_hint=key_str)
+        return result_dict
+
+    if isinstance(obj, (list, tuple, set)):
+        return [_sanitize_for_logging(x, key_hint=key_hint) for x in obj]
+
+    try:
+        if inspect.isroutine(obj) or inspect.isclass(obj):
+            module = getattr(obj, "__module__", "")
+            qn = getattr(obj, "__qualname__", getattr(obj, "__name__", "<callable>"))
+            return f"{module}.{qn}" if module else qn
+    except Exception:
+        pass
+
+    try:
+        return str(obj)
+    except Exception:
+        return "<unserializable>"
+
+
+def _trace_calls(log_path, root_dir, frame, event, arg=None):
+    if event in ["call", "return"]:
+        # Extract the filename, line number, function name, and the code object
+        filename = frame.f_code.co_filename
+        lineno = frame.f_lineno
+        func_name = frame.f_code.co_name
+        if not filename.startswith(root_dir):
+            # only log the functions in the sgl_diffusion root_dir
+            return
+        # Log every function call or return
+        try:
+            last_frame = frame.f_back
+            if last_frame is not None:
+                last_filename = last_frame.f_code.co_filename
+                last_lineno = last_frame.f_lineno
+                last_func_name = last_frame.f_code.co_name
+            else:
+                # initial frame
+                last_filename = ""
+                last_lineno = 0
+                last_func_name = ""
+            with open(log_path, "a") as f:
+                ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
+                if event == "call":
+                    f.write(
+                        f"{ts} Call to"
+                        f" {func_name} in {filename}:{lineno}"
+                        f" from {last_func_name} in {last_filename}:"
+                        f"{last_lineno}\n"
+                    )
+                else:
+                    f.write(
+                        f"{ts} Return from"
+                        f" {func_name} in {filename}:{lineno}"
+                        f" to {last_func_name} in {last_filename}:"
+                        f"{last_lineno}\n"
+                    )
+        except NameError:
+            # modules are deleted during shutdown
+            pass
+    return partial(_trace_calls, log_path, root_dir)
+
+
+def enable_trace_function_call(log_file_path: str, root_dir: str | None = None):
+    """
+    Enable tracing of every function call in code under `root_dir`.
+    This is useful for debugging hangs or crashes.
+    `log_file_path` is the path to the log file.
+    `root_dir` is the root directory of the code to trace. If None, it is the
+    sgl_diffusion root directory.
+
+    Note that this call is thread-level, any threads calling this function
+    will have the trace enabled. Other threads will not be affected.
+    """
+    logger.warning(
+        "SGLANG_DIFFUSION_TRACE_FUNCTION is enabled. It will record every"
+        " function executed by Python. This will slow down the code. It "
+        "is suggested to be used for debugging hang or crashes only."
+    )
+    logger.info("Trace frame log is saved to %s", log_file_path)
+    if root_dir is None:
+        # by default, this is the sgl_diffusion root directory
+        root_dir = os.path.dirname(os.path.dirname(__file__))
+    sys.settrace(partial(_trace_calls, log_file_path, root_dir))
+
+
+def set_uvicorn_logging_configs(server_args=None):
+    from uvicorn.config import LOGGING_CONFIG
+
+    LOGGING_CONFIG["formatters"]["default"][
+        "fmt"
+    ] = "[%(asctime)s] %(levelprefix)s %(message)s"
+    LOGGING_CONFIG["formatters"]["default"]["datefmt"] = "%Y-%m-%d %H:%M:%S"
+    LOGGING_CONFIG["formatters"]["access"][
+        "fmt"
+    ] = '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s'
+    LOGGING_CONFIG["formatters"]["access"]["datefmt"] = "%Y-%m-%d %H:%M:%S"
+
+    # Install access log path filter into LOGGING_CONFIG so it survives
+    # uvicorn's internal dictConfig() call during startup.
+    prefixes = getattr(server_args, "uvicorn_access_log_exclude_prefixes", None)
+    if prefixes:
+        _install_access_log_filter(LOGGING_CONFIG, prefixes)
+
+
+def _install_access_log_filter(config: dict, prefixes: list[str]):
+    """Register a path-based access log filter into uvicorn's LOGGING_CONFIG dict.
+
+    Only attaches to the ``access`` handler (not the ``uvicorn.access`` logger)
+    to avoid filtering the same record twice.
+    """
+    # Sanitize: drop empty strings (would match all paths) and deduplicate.
+    prefixes = [str(p) for p in prefixes if p]
+    prefixes = list(dict.fromkeys(prefixes))
+    if not prefixes:
+        return
+
+    name = "sglang_diffusion_path_filter"
+    config.setdefault("filters", {})[name] = {
+        "()": "sglang.multimodal_gen.runtime.utils.logging_utils._UvicornAccessLogFilter",
+        "prefixes": prefixes,
+    }
+
+    handler_cfg = config.get("handlers", {}).get("access")
+    if handler_cfg is not None:
+        fl = handler_cfg.setdefault("filters", [])
+        if name not in fl:
+            fl.append(name)
+
+
+class _UvicornAccessLogFilter(logging.Filter):
+    """Suppress uvicorn access logs whose path starts with an excluded prefix.
+
+    uvicorn's ``AccessFormatter`` injects ``request_line`` during ``format()``,
+    which runs *after* filters.  We therefore extract the path from
+    ``record.args`` which uvicorn populates as::
+
+        (client_addr, method, full_path, http_version, status_code)
+    """
+
+    def __init__(self, prefixes: list[str] | None = None):
+        super().__init__()
+        self.prefixes = tuple(str(p) for p in (prefixes or ()) if p)
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        args = record.args
+        if isinstance(args, tuple) and len(args) >= 3:
+            path = str(args[2]).split("?", 1)[0]
+            return not path.startswith(self.prefixes)
+        return True
+
+
+def configure_logger(server_args, prefix: str = ""):
+    log_format = f"[%(asctime)s{prefix}] %(message)s"
+    datefmt = "%m-%d %H:%M:%S"
+
+    formatter = ColoredFormatter(log_format, datefmt=datefmt)
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setFormatter(formatter)
+
+    root = logging.getLogger()
+    root.handlers.clear()
+    root.addHandler(handler)
+    root.setLevel(getattr(logging, server_args.log_level.upper()))
+
+    set_uvicorn_logging_configs(server_args)
+
+
+@lru_cache(maxsize=1)
+def get_log_level() -> int:
+    root = logging.getLogger()
+    return root.level
+
+
+def suppress_loggers(loggers_to_suppress: list[str], level: int = logging.WARNING):
+    original_levels = {}
+
+    for logger_name in loggers_to_suppress:
+        logger = logging.getLogger(logger_name)
+        original_levels[logger_name] = logger.level
+        logger.setLevel(level)
+
+    return original_levels
+
+
+def globally_suppress_loggers():
+    # globally suppress some obsessive loggers
+    target_names = [
+        "imageio",
+        "imageio_ffmpeg",
+        "PIL",
+        "PIL_Image",
+        "python_multipart.multipart",
+        "filelock",
+        "urllib3",
+        "httpx",
+        "httpcore",
+        "diffusers.quantizers.torchao.torchao_quantizer",
+        "transformers.processing_utils",
+        "flash_attn.cute.cache_utils",
+    ]
+
+    for name in target_names:
+        logging.getLogger(name).setLevel(logging.ERROR)
+
+
+# source: https://github.com/vllm-project/vllm/blob/a11f4a81e027efd9ef783b943489c222950ac989/vllm/utils/system_utils.py#L60
+@contextlib.contextmanager
+def suppress_stdout():
+    """
+    Suppress stdout from C libraries at the file descriptor level.
+
+    Only suppresses stdout, not stderr, to preserve error messages.
+    Example:
+        with suppress_stdout():
+            # C library calls that would normally print to stdout
+            torch.distributed.new_group(ranks, backend="gloo")
+    """
+    # Don't suppress if logging level is DEBUG
+
+    stdout_fd = sys.stdout.fileno()
+    stdout_dup = os.dup(stdout_fd)
+    devnull_fd = os.open(os.devnull, os.O_WRONLY)
+
+    try:
+        sys.stdout.flush()
+        os.dup2(devnull_fd, stdout_fd)
+        yield
+    finally:
+        sys.stdout.flush()
+        os.dup2(stdout_dup, stdout_fd)
+        os.close(stdout_dup)
+        os.close(devnull_fd)
+
+
+class GenerationTimer:
+    def __init__(self):
+        self.start_time = 0.0
+        self.end_time = 0.0
+        self.duration = 0.0
+
+
+@contextmanager
+def log_generation_timer(
+    logger: logging.Logger,
+    prompt: str,
+    request_idx: int | None = None,
+    total_requests: int | None = None,
+):
+    if request_idx is not None and total_requests is not None:
+        logger.info(
+            "Processing prompt %d/%d: %s",
+            request_idx,
+            total_requests,
+            _sanitize_for_logging(prompt, key_hint="prompt"),
+        )
+
+    timer = GenerationTimer()
+    timer.start_time = time.perf_counter()
+    try:
+        yield timer
+        timer.end_time = time.perf_counter()
+        timer.duration = timer.end_time - timer.start_time
+        logger.info(
+            f"Pixel data generated successfully in {GREEN}%.2f{RESET} seconds",
+            timer.duration,
+        )
+    except Exception as e:
+        if request_idx is not None:
+            logger.error(
+                "Failed to generate output for prompt %d: %s",
+                request_idx,
+                e,
+                exc_info=True,
+            )
+        else:
+            logger.error(
+                f"Failed to generate output for prompt: {e}",
+                exc_info=True,
+            )
+        raise
+
+
+def log_batch_completion(
+    logger: logging.Logger, num_outputs: int, total_time: float
+) -> None:
+    logger.info(
+        f"Completed batch processing. Generated %d outputs in {GREEN}%.2f{RESET} seconds",
+        num_outputs,
+        total_time,
+    )

python/sglang/multimodal_gen/runtime/utils/model_overlay.py

@@ -0,0 +1,780 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import glob
+import hashlib
+import importlib.util
+import json
+import os
+import shutil
+from typing import Any, Callable, cast
+
+from huggingface_hub.errors import (
+    LocalEntryNotFoundError,
+    RepositoryNotFoundError,
+    RevisionNotFoundError,
+)
+from requests.exceptions import ConnectionError as RequestsConnectionError
+from requests.exceptions import RequestException
+
+from sglang.multimodal_gen.runtime.loader.weight_utils import get_lock
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.utils import load_diffusion_overlay_registry_from_env
+
+logger = init_logger(__name__)
+
+# Built-in diffusion model overlay registry.
+BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = {
+    "Lightricks/LTX-2.3": {
+        "overlay_repo_id": "MickJ/LTX-2.3-overlay",
+        "overlay_revision": "e0cc94f279ec16bb87c230134d40319f6ce40c5e",
+    },
+    "jdopensource/JoyAI-Echo": {
+        "overlay_repo_id": "Niehen6174/JoyAI-Echo-overlay",
+        "overlay_revision": "0a19f315c96532b7a5f61bcd765d1fefdd83dc7d",
+    },
+    "Efficient-Large-Model/SANA-WM_bidirectional": {
+        "overlay_repo_id": "sjmshsh/SANA-WM_bidirectional-overlay",
+        "overlay_revision": "e611beacbcc0cf33c676306ae0eb89f149e044ad",
+    },
+    "Efficient-Large-Model/SANA-WM_streaming": {
+        "overlay_repo_id": "AgainstEntropy/SANA-WM_streaming-overlay",
+        "overlay_revision": "62c6840871ecc3559189047513ba0670e1bf62e7",
+    },
+}
+
+
+MODEL_OVERLAY_METADATA_PATTERNS = [
+    "*.json",
+    "*.md",
+    "*.py",
+    "*.txt",
+    "**/*.json",
+    "**/*.md",
+    "**/*.py",
+    "**/*.txt",
+]
+
+_MATERIALIZED_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pth", ".pt")
+_MATERIALIZED_CONFIG_ONLY_COMPONENTS = {
+    "feature_extractor",
+    "image_processor",
+    "processor",
+    "scheduler",
+    "tokenizer",
+    "tokenizer_2",
+}
+
+_MODEL_OVERLAY_REGISTRY_CACHE: dict[str, dict[str, Any]] | None = None
+
+
+def _compute_overlay_fingerprint(overlay_dir: str) -> str:
+    hasher = hashlib.sha256()
+    for root, dir_names, file_names in os.walk(overlay_dir):
+        dir_names[:] = sorted(
+            d for d in dir_names if d != "__pycache__" and not d.endswith(".egg-info")
+        )
+        for file_name in sorted(file_names):
+            if file_name.endswith((".safetensors", ".bin", ".pth", ".pt")):
+                continue
+            file_path = os.path.join(root, file_name)
+            rel_path = os.path.relpath(file_path, overlay_dir).replace(os.sep, "/")
+            hasher.update(rel_path.encode("utf-8"))
+            with open(file_path, "rb") as f:
+                hasher.update(hashlib.sha256(f.read()).digest())
+    return hasher.hexdigest()
+
+
+def _resolve_bundled_overlay_dir(overlay_spec: dict[str, Any]) -> str | None:
+    bundled_overlay_subdir = overlay_spec.get("bundled_overlay_subdir")
+    if not bundled_overlay_subdir:
+        return None
+    bundled_overlay_dir = os.path.abspath(
+        os.path.join(
+            os.path.dirname(__file__),
+            "..",
+            "..",
+            "model_overlays",
+            str(bundled_overlay_subdir),
+        )
+    )
+    if not os.path.isdir(bundled_overlay_dir):
+        return None
+    if load_overlay_manifest_if_present(bundled_overlay_dir) is None:
+        return None
+    return bundled_overlay_dir
+
+
+def get_diffusion_cache_root() -> str:
+    return os.path.expanduser(
+        os.getenv("SGLANG_DIFFUSION_CACHE_ROOT", "~/.cache/sgl_diffusion")
+    )
+
+
+def clear_model_overlay_registry_cache() -> None:
+    global _MODEL_OVERLAY_REGISTRY_CACHE
+    _MODEL_OVERLAY_REGISTRY_CACHE = None
+
+
+def _load_model_overlay_registry() -> dict[str, dict[str, Any]]:
+    global _MODEL_OVERLAY_REGISTRY_CACHE
+    if _MODEL_OVERLAY_REGISTRY_CACHE is not None:
+        return _MODEL_OVERLAY_REGISTRY_CACHE
+
+    # Built-in registry is the stable default path; env only overrides it.
+    normalized = _normalize_model_overlay_registry(BUILTIN_MODEL_OVERLAY_REGISTRY)
+
+    env_registry = load_diffusion_overlay_registry_from_env()
+    if not env_registry:
+        _MODEL_OVERLAY_REGISTRY_CACHE = normalized
+        return _MODEL_OVERLAY_REGISTRY_CACHE
+
+    normalized.update(_normalize_model_overlay_registry(env_registry))
+    _MODEL_OVERLAY_REGISTRY_CACHE = normalized
+    return _MODEL_OVERLAY_REGISTRY_CACHE
+
+
+def _normalize_model_overlay_registry(
+    payload: dict[str, Any],
+) -> dict[str, dict[str, Any]]:
+    normalized: dict[str, dict[str, Any]] = {}
+    for source_model_id, spec in payload.items():
+        if isinstance(spec, str):
+            normalized[source_model_id] = {"overlay_repo_id": spec}
+            continue
+        if not isinstance(spec, dict):
+            raise ValueError(
+                "Overlay registry values must be either strings or JSON objects"
+            )
+        overlay_repo_id = spec.get("overlay_repo_id")
+        if not overlay_repo_id:
+            raise ValueError(
+                f"Overlay registry entry for {source_model_id!r} is missing overlay_repo_id"
+            )
+        normalized[source_model_id] = dict(spec)
+    return normalized
+
+
+def resolve_model_overlay(model_name_or_path: str) -> dict[str, Any] | None:
+    registry = _load_model_overlay_registry()
+    return registry.get(model_name_or_path)
+
+
+def resolve_model_overlay_target(
+    model_name_or_path: str,
+) -> tuple[str, dict[str, Any]] | None:
+    registry = _load_model_overlay_registry()
+
+    exact = registry.get(model_name_or_path)
+    if exact is not None:
+        return model_name_or_path, exact
+
+    if os.path.exists(model_name_or_path):
+        # Local source dirs do not have a repo id, so match them by basename.
+        base_name = os.path.basename(os.path.normpath(model_name_or_path))
+        normalized_path = (
+            os.path.normpath(model_name_or_path).lower().replace(os.sep, "/")
+        )
+        for source_model_id, spec in registry.items():
+            if base_name == source_model_id.rsplit("/", 1)[-1]:
+                return source_model_id, spec
+            cache_repo_fragment = (
+                f"models--{source_model_id.lower().replace('/', '--')}"
+            )
+            if cache_repo_fragment in normalized_path:
+                return source_model_id, spec
+
+    return None
+
+
+def load_overlay_manifest_if_present(overlay_dir: str) -> dict[str, Any] | None:
+    overlay_manifest_path = os.path.join(
+        overlay_dir, "_overlay", "overlay_manifest.json"
+    )
+    if not os.path.exists(overlay_manifest_path):
+        return None
+    with open(overlay_manifest_path, encoding="utf-8") as f:
+        manifest = cast(dict[str, Any], json.load(f))
+    return manifest
+
+
+def load_model_index_from_dir(model_dir: str) -> dict[str, Any]:
+    model_index_path = os.path.join(model_dir, "model_index.json")
+    if not os.path.exists(model_index_path):
+        raise ValueError(f"model_index.json not found under {model_dir}")
+    with open(model_index_path, encoding="utf-8") as f:
+        config = cast(dict[str, Any], json.load(f))
+    if "_class_name" not in config or "_diffusers_version" not in config:
+        raise ValueError(f"Invalid model_index.json under {model_dir}")
+    config["pipeline_name"] = config["_class_name"]
+    return config
+
+
+def _component_has_weight_file(component_dir: str) -> bool:
+    for root, _, file_names in os.walk(component_dir):
+        if any(
+            file_name.endswith(_MATERIALIZED_WEIGHT_SUFFIXES)
+            and os.path.isfile(os.path.join(root, file_name))
+            for file_name in file_names
+        ):
+            return True
+    return False
+
+
+def _materialized_overlay_has_component_weights(model_dir: str) -> bool:
+    model_index = load_model_index_from_dir(model_dir)
+    for component_name, entry in model_index.items():
+        if (
+            component_name.startswith("_")
+            or component_name == "pipeline_name"
+            or component_name in _MATERIALIZED_CONFIG_ONLY_COMPONENTS
+            or not isinstance(entry, list)
+        ):
+            continue
+        component_dir = os.path.join(model_dir, component_name)
+        if not os.path.isdir(component_dir) or not _component_has_weight_file(
+            component_dir
+        ):
+            logger.warning(
+                "Materialized overlay cache for %s is missing weights for component %s",
+                model_dir,
+                component_name,
+            )
+            return False
+    return True
+
+
+def _materialized_overlay_cache_complete(
+    final_dir: str,
+    marker_path: str,
+    verify_diffusers_model_complete_fn: Callable[[str], bool],
+) -> bool:
+    return (
+        verify_diffusers_model_complete_fn(final_dir)
+        and os.path.exists(marker_path)
+        and _materialized_overlay_has_component_weights(final_dir)
+    )
+
+
+def _ensure_dir(path: str) -> None:
+    os.makedirs(path, exist_ok=True)
+
+
+def _find_missing_required_paths(
+    root_dir: str, required_paths: list[str] | tuple[str, ...]
+) -> list[str]:
+    missing: list[str] = []
+    for rel_path in required_paths:
+        if not os.path.exists(os.path.join(root_dir, rel_path)):
+            missing.append(rel_path)
+    return missing
+
+
+def _link_or_copy_file(src: str, dst: str) -> None:
+    src = os.path.realpath(src)
+    _ensure_dir(os.path.dirname(dst))
+    if os.path.lexists(dst):
+        os.remove(dst)
+    try:
+        os.link(src, dst)
+        return
+    except OSError:
+        pass
+    try:
+        os.symlink(src, dst)
+        return
+    except OSError:
+        pass
+    shutil.copy2(src, dst)
+
+
+def _copytree_link_or_copy(src_dir: str, dst_dir: str) -> None:
+    for root, _, files in os.walk(src_dir):
+        rel_root = os.path.relpath(root, src_dir)
+        target_root = dst_dir if rel_root == "." else os.path.join(dst_dir, rel_root)
+        _ensure_dir(target_root)
+        for file_name in files:
+            src_file = os.path.join(root, file_name)
+            dst_file = os.path.join(target_root, file_name)
+            _link_or_copy_file(src_file, dst_file)
+
+
+def ensure_overlay_source_dir_complete(
+    *,
+    source_model_id: str,
+    source_dir: str,
+    manifest: dict[str, Any],
+    local_dir: str | None,
+    allow_patterns: list[str] | None,
+    download: bool,
+    snapshot_download_fn: Callable[..., str],
+) -> str:
+    required_source_files = cast(
+        list[str], list(manifest.get("required_source_files", []))
+    )
+    if not required_source_files:
+        return source_dir
+
+    # Metadata-only overlays often need a partial source snapshot. Re-download
+    # only when the current source dir is missing required files.
+    missing_paths = _find_missing_required_paths(source_dir, required_source_files)
+    if not missing_paths:
+        return source_dir
+
+    if not download:
+        raise ValueError(
+            f"Overlay source model {source_model_id} is missing required files "
+            f"{missing_paths} and download=False."
+        )
+
+    logger.warning(
+        "Overlay source model %s is missing required files %s. "
+        "Re-downloading source snapshot.",
+        source_model_id,
+        missing_paths,
+    )
+    source_allow_patterns = manifest.get("source_allow_patterns")
+    effective_allow_patterns = (
+        cast(list[str] | None, source_allow_patterns)
+        if source_allow_patterns is not None
+        else allow_patterns
+    )
+    with get_lock(source_model_id).acquire(poll_interval=2):
+        source_dir = snapshot_download_fn(
+            repo_id=source_model_id,
+            ignore_patterns=["*.onnx", "*.msgpack"],
+            allow_patterns=effective_allow_patterns,
+            local_dir=local_dir,
+            max_workers=8,
+            force_download=True,
+        )
+    missing_after_redownload = _find_missing_required_paths(
+        source_dir, required_source_files
+    )
+    if missing_after_redownload:
+        raise ValueError(
+            f"Overlay source model {source_model_id} is still missing required files "
+            f"{missing_after_redownload} after re-download."
+        )
+    return str(source_dir)
+
+
+def resolve_direct_overlay_repo(
+    model_name_or_path: str,
+    *,
+    hf_hub_download_fn: Callable[..., str],
+) -> tuple[dict[str, Any], str, dict[str, Any]] | None:
+    if os.path.exists(model_name_or_path):
+        manifest = load_overlay_manifest_if_present(model_name_or_path)
+        if manifest is None:
+            return None
+        source_model_id = manifest.get("source_model_id")
+        if not source_model_id:
+            raise ValueError(
+                f"Overlay repo {model_name_or_path} is missing source_model_id in _overlay/overlay_manifest.json"
+            )
+        overlay_spec = {
+            "overlay_repo_id": model_name_or_path,
+            "overlay_revision": "local",
+        }
+        return overlay_spec, model_name_or_path, manifest
+
+    try:
+        manifest_path = hf_hub_download_fn(
+            repo_id=model_name_or_path,
+            filename="_overlay/overlay_manifest.json",
+        )
+        overlay_dir = os.path.dirname(os.path.dirname(manifest_path))
+    except (
+        RepositoryNotFoundError,
+        RevisionNotFoundError,
+        LocalEntryNotFoundError,
+        RequestsConnectionError,
+        RequestException,
+    ):
+        return None
+    except Exception:
+        return None
+
+    manifest = load_overlay_manifest_if_present(overlay_dir)
+    if manifest is None:
+        return None
+    source_model_id = manifest.get("source_model_id")
+    if not source_model_id:
+        raise ValueError(
+            f"Overlay repo {model_name_or_path} is missing source_model_id in _overlay/overlay_manifest.json"
+        )
+    overlay_spec = {
+        "overlay_repo_id": model_name_or_path,
+        "overlay_revision": "main",
+    }
+    return overlay_spec, overlay_dir, manifest
+
+
+def download_overlay_metadata(
+    source_model_id: str,
+    overlay_spec: dict[str, Any],
+    *,
+    snapshot_download_fn: Callable[..., str],
+) -> str:
+    bundled_overlay_dir = _resolve_bundled_overlay_dir(overlay_spec)
+    if bundled_overlay_dir is not None:
+        logger.info(
+            "Using bundled overlay metadata for %s from %s",
+            source_model_id,
+            bundled_overlay_dir,
+        )
+        return bundled_overlay_dir
+
+    overlay_repo_id = str(overlay_spec["overlay_repo_id"])
+    if os.path.exists(overlay_repo_id):
+        logger.info(
+            "Using local overlay metadata for %s from %s",
+            source_model_id,
+            overlay_repo_id,
+        )
+        return overlay_repo_id
+    revision = overlay_spec.get("overlay_revision")
+    logger.info(
+        "Downloading overlay metadata for %s from %s",
+        source_model_id,
+        overlay_repo_id,
+    )
+    return str(
+        snapshot_download_fn(
+            repo_id=overlay_repo_id,
+            allow_patterns=MODEL_OVERLAY_METADATA_PATTERNS,
+            revision=revision,
+            max_workers=4,
+        )
+    )
+
+
+def _apply_overlay_file_mappings(
+    *,
+    source_dir: str,
+    output_dir: str,
+    file_mappings: list[dict[str, Any]],
+) -> None:
+    for mapping in file_mappings:
+        mapping_type = mapping.get("type", "file")
+        src_rel = mapping.get("src")
+        if not src_rel:
+            raise ValueError(f"Overlay file mapping is missing src: {mapping}")
+        src_path = os.path.join(source_dir, src_rel)
+        if mapping_type == "tree":
+            if not os.path.isdir(src_path):
+                raise ValueError(f"Tree mapping source does not exist: {src_path}")
+            dst_dir = os.path.join(output_dir, str(mapping.get("dst_dir", src_rel)))
+            _copytree_link_or_copy(src_path, dst_dir)
+            continue
+        if mapping_type == "glob":
+            matched = glob.glob(src_path, recursive=True)
+            if not matched:
+                raise ValueError(f"Glob mapping matched no files: {src_path}")
+            for matched_path in matched:
+                if os.path.isdir(matched_path):
+                    continue
+                rel_path = os.path.relpath(matched_path, source_dir)
+                dst_path = os.path.join(output_dir, rel_path)
+                _link_or_copy_file(matched_path, dst_path)
+            continue
+
+        if not os.path.isfile(src_path):
+            raise ValueError(f"File mapping source does not exist: {src_path}")
+        dst_rel = str(mapping.get("dst", os.path.basename(src_rel)))
+        dst_path = os.path.join(output_dir, dst_rel)
+        _link_or_copy_file(src_path, dst_path)
+
+
+def _run_overlay_custom_materializer(
+    *,
+    overlay_dir: str,
+    source_dir: str,
+    output_dir: str,
+    manifest: dict[str, Any],
+) -> None:
+    custom_materializer = manifest.get("custom_materializer")
+    if not custom_materializer:
+        return
+    script_path = os.path.join(overlay_dir, str(custom_materializer))
+    if not os.path.exists(script_path):
+        raise ValueError(f"Custom materializer script not found: {script_path}")
+
+    spec = importlib.util.spec_from_file_location(
+        "_sglang_overlay_materializer", script_path
+    )
+    if spec is None or spec.loader is None:
+        raise ValueError(f"Failed to import custom materializer: {script_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    materialize_fn = getattr(module, "materialize", None)
+    if materialize_fn is None:
+        raise ValueError(
+            f"Custom materializer {script_path} must define materialize(...)"
+        )
+
+    materialize_fn(
+        overlay_dir=overlay_dir,
+        source_dir=source_dir,
+        output_dir=output_dir,
+        manifest=manifest,
+    )
+
+
+def materialize_overlay_model(
+    *,
+    source_model_id: str,
+    overlay_spec: dict[str, Any],
+    overlay_dir: str,
+    source_dir: str,
+    verify_diffusers_model_complete_fn: Callable[[str], bool],
+) -> str:
+    overlay_manifest_path = os.path.join(
+        overlay_dir, "_overlay", "overlay_manifest.json"
+    )
+    if not os.path.exists(overlay_manifest_path):
+        raise ValueError(
+            f"Overlay repo for {source_model_id} is missing _overlay/overlay_manifest.json"
+        )
+
+    with open(overlay_manifest_path, encoding="utf-8") as f:
+        manifest = cast(dict[str, Any], json.load(f))
+
+    materializer_version = str(manifest.get("materializer_version", "v1"))
+    overlay_repo_id = str(overlay_spec["overlay_repo_id"])
+    overlay_revision = str(overlay_spec.get("overlay_revision", "main"))
+    overlay_fingerprint = _compute_overlay_fingerprint(overlay_dir)
+    cache_key = hashlib.sha256(
+        json.dumps(
+            {
+                "source_model_id": source_model_id,
+                "overlay_repo_id": overlay_repo_id,
+                "overlay_revision": overlay_revision,
+                "materializer_version": materializer_version,
+                "overlay_fingerprint": overlay_fingerprint,
+            },
+            sort_keys=True,
+        ).encode("utf-8")
+    ).hexdigest()[:16]
+    cache_root = os.path.join(get_diffusion_cache_root(), "materialized_models")
+    _ensure_dir(cache_root)
+    safe_name = source_model_id.replace("/", "__")
+    final_dir = os.path.join(cache_root, f"{safe_name}-{cache_key}")
+    marker_path = os.path.join(final_dir, ".sglang_overlay_materialized.json")
+    if _materialized_overlay_cache_complete(
+        final_dir, marker_path, verify_diffusers_model_complete_fn
+    ):
+        return final_dir
+
+    lock_name = (
+        f"overlay-materialize::{source_model_id}::{overlay_repo_id}::{overlay_revision}"
+    )
+    with get_lock(lock_name).acquire(poll_interval=2):
+        if _materialized_overlay_cache_complete(
+            final_dir, marker_path, verify_diffusers_model_complete_fn
+        ):
+            return final_dir
+
+        logger.info(
+            "Materializing overlay model for %s into %s",
+            source_model_id,
+            final_dir,
+        )
+        logger.info(
+            "Overlay source repo: %s, overlay repo: %s@%s",
+            source_model_id,
+            overlay_repo_id,
+            overlay_revision,
+        )
+        tmp_dir = final_dir + ".tmp"
+        if os.path.exists(tmp_dir):
+            shutil.rmtree(tmp_dir)
+        if os.path.exists(final_dir):
+            shutil.rmtree(final_dir)
+        logger.info("Copying overlay metadata into temporary materialized directory")
+        shutil.copytree(
+            overlay_dir,
+            tmp_dir,
+            ignore=shutil.ignore_patterns("*.safetensors", "*.bin", "*.pth", "*.pt"),
+        )
+
+        overlay_hidden_dir = os.path.join(tmp_dir, "_overlay")
+        if os.path.isdir(overlay_hidden_dir):
+            shutil.rmtree(overlay_hidden_dir)
+
+        file_mappings = manifest.get("file_mappings", [])
+        if file_mappings:
+            logger.info("Applying %d overlay file mappings", len(file_mappings))
+            _apply_overlay_file_mappings(
+                source_dir=source_dir,
+                output_dir=tmp_dir,
+                file_mappings=cast(list[dict[str, Any]], file_mappings),
+            )
+        if manifest.get("custom_materializer"):
+            logger.info(
+                "Running custom overlay materializer: %s",
+                manifest["custom_materializer"],
+            )
+        _run_overlay_custom_materializer(
+            overlay_dir=overlay_dir,
+            source_dir=source_dir,
+            output_dir=tmp_dir,
+            manifest=manifest,
+        )
+
+        with open(marker_path.replace(final_dir, tmp_dir), "w", encoding="utf-8") as f:
+            json.dump(
+                {
+                    "source_model_id": source_model_id,
+                    "source_dir": source_dir,
+                    "overlay_repo_id": overlay_repo_id,
+                    "overlay_revision": overlay_revision,
+                    "materializer_version": materializer_version,
+                    "overlay_fingerprint": overlay_fingerprint,
+                },
+                f,
+                indent=2,
+                sort_keys=True,
+            )
+
+        os.replace(tmp_dir, final_dir)
+        logger.info("Overlay materialization finished: %s", final_dir)
+
+    return final_dir
+
+
+def maybe_load_overlay_model_index(
+    model_name_or_path: str,
+    *,
+    snapshot_download_fn: Callable[..., str],
+    hf_hub_download_fn: Callable[..., str],
+) -> dict[str, Any] | None:
+    if os.path.exists(model_name_or_path):
+        # A local overlay repo already contains the model_index we need.
+        if load_overlay_manifest_if_present(model_name_or_path) is not None:
+            return load_model_index_from_dir(model_name_or_path)
+
+    overlay_target = resolve_model_overlay_target(model_name_or_path)
+    if overlay_target is not None:
+        # Registry-mapped source model ids first resolve to overlay metadata.
+        source_model_id, overlay_spec = overlay_target
+        overlay_dir = download_overlay_metadata(
+            source_model_id,
+            overlay_spec,
+            snapshot_download_fn=snapshot_download_fn,
+        )
+        return load_model_index_from_dir(overlay_dir)
+
+    direct_overlay = resolve_direct_overlay_repo(
+        model_name_or_path, hf_hub_download_fn=hf_hub_download_fn
+    )
+    if direct_overlay is None:
+        return None
+
+    _, overlay_dir, _ = direct_overlay
+    return load_model_index_from_dir(overlay_dir)
+
+
+def maybe_resolve_overlay_model_path(
+    model_name_or_path: str,
+    *,
+    local_dir: str | None,
+    download: bool,
+    allow_patterns: list[str] | None,
+    snapshot_download_fn: Callable[..., str],
+    hf_hub_download_fn: Callable[..., str],
+    verify_diffusers_model_complete_fn: Callable[[str], bool],
+    base_model_download_fn: Callable[..., str],
+) -> str | None:
+    overlay_target = resolve_model_overlay_target(model_name_or_path)
+    if overlay_target is not None:
+        source_model_id, overlay_spec = overlay_target
+        overlay_dir = download_overlay_metadata(
+            source_model_id,
+            overlay_spec,
+            snapshot_download_fn=snapshot_download_fn,
+        )
+        manifest = load_overlay_manifest_if_present(overlay_dir)
+        if manifest is None:
+            # Full diffusers overlays do not need materialization.
+            return base_model_download_fn(
+                str(overlay_spec["overlay_repo_id"]),
+                local_dir=local_dir,
+                download=download,
+                allow_patterns=allow_patterns,
+                force_diffusers_model=True,
+                skip_overlay_resolution=True,
+            )
+        source_allow_patterns = cast(
+            list[str] | None, manifest.get("source_allow_patterns")
+        )
+        # For local source paths, reuse the directory directly instead of
+        # round-tripping through snapshot_download.
+        source_dir = (
+            model_name_or_path
+            if os.path.exists(model_name_or_path)
+            else base_model_download_fn(
+                source_model_id,
+                local_dir=local_dir,
+                download=download,
+                allow_patterns=source_allow_patterns or allow_patterns,
+                force_diffusers_model=False,
+                skip_overlay_resolution=True,
+            )
+        )
+        source_dir = ensure_overlay_source_dir_complete(
+            source_model_id=source_model_id,
+            source_dir=source_dir,
+            manifest=manifest,
+            local_dir=local_dir,
+            allow_patterns=allow_patterns,
+            download=download,
+            snapshot_download_fn=snapshot_download_fn,
+        )
+        return materialize_overlay_model(
+            source_model_id=source_model_id,
+            overlay_spec=overlay_spec,
+            overlay_dir=overlay_dir,
+            source_dir=source_dir,
+            verify_diffusers_model_complete_fn=verify_diffusers_model_complete_fn,
+        )
+
+    direct_overlay = resolve_direct_overlay_repo(
+        model_name_or_path, hf_hub_download_fn=hf_hub_download_fn
+    )
+    if direct_overlay is None:
+        return None
+
+    overlay_spec, overlay_dir, manifest = direct_overlay
+    source_model_id = str(manifest["source_model_id"])
+    # Direct overlay repos are always metadata-only; they need the original
+    # source weights before they can be materialized into a diffusers-like dir.
+    source_allow_patterns = cast(
+        list[str] | None, manifest.get("source_allow_patterns")
+    )
+    source_dir = base_model_download_fn(
+        source_model_id,
+        local_dir=local_dir,
+        download=download,
+        allow_patterns=source_allow_patterns or allow_patterns,
+        force_diffusers_model=False,
+        skip_overlay_resolution=True,
+    )
+    source_dir = ensure_overlay_source_dir_complete(
+        source_model_id=source_model_id,
+        source_dir=source_dir,
+        manifest=manifest,
+        local_dir=local_dir,
+        allow_patterns=allow_patterns,
+        download=download,
+        snapshot_download_fn=snapshot_download_fn,
+    )
+    return materialize_overlay_model(
+        source_model_id=source_model_id,
+        overlay_spec=overlay_spec,
+        overlay_dir=overlay_dir,
+        source_dir=source_dir,
+        verify_diffusers_model_complete_fn=verify_diffusers_model_complete_fn,
+    )

python/sglang/multimodal_gen/runtime/utils/nvtx_pytorch_hooks.py

@@ -0,0 +1,214 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""PyTorch hooks for layerwise NVTX profiling in SGLang Diffusion.
+
+Mirrors the structure of ``sglang.srt.utils.nvtx_pytorch_hooks.PytHooks``
+but uses a compact ``{name} in={shapes}`` marker format that is well-suited
+to DiT transformer blocks. See
+``sglang.srt.utils.nvtx_pytorch_hooks`` for the LLM-runtime equivalent
+that emits a richer per-layer parameter dict.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Iterator
+from typing import Any
+
+import torch
+import torch.cuda.nvtx as nvtx
+from torch.utils.hooks import RemovableHandle
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# Module types that are too lightweight to warrant their own NVTX range.
+# Skipping them keeps the captured timeline readable.
+_DEFAULT_SKIP_TYPES: tuple[type, ...] = (
+    torch.nn.Identity,
+    torch.nn.Dropout,
+    torch.nn.Dropout1d,
+    torch.nn.Dropout2d,
+    torch.nn.Dropout3d,
+)
+
+
+@contextlib.contextmanager
+def maybe_nvtx_range(name: str, enabled: bool = True) -> Iterator[None]:
+    """Context manager that wraps a block of work in an NVTX range.
+
+    Calls ``range_push`` / ``range_pop`` directly rather than going through
+    :func:`torch.cuda.nvtx.range`, which would otherwise interpret ``name`` as a
+    ``str.format`` template (so a literal ``{`` in the marker would raise
+    ``KeyError``). The ``range_pop`` is invoked from the ``finally`` clause, so
+    exceptions raised inside the ``with`` block cannot leak a half-open range.
+
+    When ``enabled`` is ``False`` the function is a zero-cost no-op, suitable
+    for use under a per-request gate (e.g. warmup exclusion).
+    """
+    if not enabled:
+        yield
+        return
+    nvtx.range_push(name)
+    try:
+        yield
+    finally:
+        nvtx.range_pop()
+
+
+class DiffusionNvtxHooks:
+    """Register NVTX markers around each submodule forward pass.
+
+    Each registered module emits an NVTX range covering its forward pass.
+    The range name encodes the qualified module name and the input tensor
+    shapes for downstream identification in Nsight Systems.
+
+    Hook handles are retained so they can be removed via :meth:`remove_hooks`;
+    the same instance must not be reused across multiple model instances.
+    """
+
+    def __init__(self, skip_types: tuple[type, ...] = _DEFAULT_SKIP_TYPES) -> None:
+        self._skip_types = skip_types
+        self._module_to_name_map: dict[torch.nn.Module, str] = {}
+        self._hook_handles: list[RemovableHandle] = []
+        # Caller must explicitly enable via ``set_enabled``. Default off
+        # so a forward that bypasses the component-use gate (e.g. an early
+        # warmup pass) cannot accidentally pollute the captured timeline.
+        self._enabled: bool = False
+
+    def register_hooks(
+        self,
+        model: torch.nn.Module,
+        prefix: str = "",
+    ) -> int:
+        """Walk ``model`` and attach forward pre/post hooks to every module.
+
+        Args:
+            model: Root module to instrument.
+            prefix: Optional name prefix prepended to every emitted range.
+
+        Returns:
+            Number of modules instrumented.
+
+        Notes:
+            Weight-tied or otherwise duplicated module instances are
+            skipped (the first occurrence wins) so each forward pass
+            produces exactly one NVTX range.
+        """
+        instrumented = 0
+        for name, module in model.named_modules(prefix=prefix):
+            if isinstance(module, self._skip_types):
+                continue
+            # Skip duplicate module instances (e.g., weight-tied layers).
+            # The check must happen before hook registration to avoid
+            # double-emitting NVTX ranges on the second occurrence.
+            if module in self._module_to_name_map:
+                logger.debug(
+                    "NVTX: module %s already registered as '%s', skipping '%s'",
+                    type(module).__name__,
+                    self._module_to_name_map[module],
+                    name,
+                )
+                continue
+            self._module_to_name_map[module] = name
+            self._hook_handles.append(
+                module.register_forward_pre_hook(
+                    self._forward_pre_hook, with_kwargs=True
+                )
+            )
+            # ``always_call=True`` (PyTorch 2.0+) guarantees the post-hook
+            # still fires when ``forward`` raises, so an OOM or assertion
+            # inside the wrapped module cannot leak a half-open NVTX range.
+            self._hook_handles.append(
+                module.register_forward_hook(self._forward_hook, always_call=True)
+            )
+            instrumented += 1
+        return instrumented
+
+    def remove_hooks(self) -> None:
+        """Remove every hook registered by this instance.
+
+        Safe to call multiple times; subsequent calls are no-ops. The
+        bookkeeping is cleared in a ``finally`` so a misbehaving
+        ``handle.remove()`` cannot leave the instance with stale
+        handles or name-map entries.
+        """
+        try:
+            for handle in self._hook_handles:
+                handle.remove()
+        finally:
+            self._hook_handles.clear()
+            self._module_to_name_map.clear()
+
+    def set_enabled(self, enabled: bool) -> None:
+        """Toggle whether the registered hooks emit NVTX ranges.
+
+        When disabled, both the pre- and post-hooks early-return, so each
+        forward produces a matched (push, pop) pair of "no-ops" — no range
+        leak and no half-open range across the toggle.
+        """
+        self._enabled = enabled
+
+    # ------------------------------------------------------------------ hooks
+
+    def _forward_pre_hook(
+        self,
+        module: torch.nn.Module,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> None:
+        if not self._enabled:
+            return
+        name = self._module_to_name_map.get(module, "unknown")
+        shapes = _collect_input_shapes(args, kwargs)
+        marker = f"{name} in={shapes}" if shapes else name
+        nvtx.range_push(marker)
+
+    def _forward_hook(
+        self,
+        module: torch.nn.Module,
+        _args: Any,
+        _output: Any,
+    ) -> None:
+        if not self._enabled:
+            return
+        nvtx.range_pop()
+
+
+def _collect_input_shapes(
+    args: tuple[Any, ...], kwargs: dict[str, Any] | None = None
+) -> list[list[int]]:
+    """Best-effort extraction of input tensor shapes for marker labels.
+
+    Walks positional ``args`` and keyword ``kwargs`` values, recursing into
+    lists and tuples (so DiT inputs like ``image_rotary_emb=(cos, sin)`` are
+    captured). Non-tensor scalars, ``None``, dicts, and arbitrary objects are
+    silently skipped.
+    """
+    shapes: list[list[int]] = []
+    _append_tensor_shapes(args, shapes)
+    if kwargs:
+        _append_tensor_shapes(tuple(kwargs.values()), shapes)
+    return shapes
+
+
+def _append_tensor_shapes(items: Any, shapes: list[list[int]]) -> None:
+    if isinstance(items, torch.Tensor):
+        shapes.append(list(items.size()))
+        return
+    if isinstance(items, (list, tuple)):
+        for item in items:
+            _append_tensor_shapes(item, shapes)

python/sglang/multimodal_gen/runtime/utils/perf_logger.py

@@ -0,0 +1,382 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+import dataclasses
+import json
+import logging
+import os
+import subprocess
+import sys
+import time
+from datetime import datetime
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import torch
+from dateutil.tz import UTC
+
+import sglang
+import sglang.multimodal_gen.envs as envs
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import (
+    CYAN,
+    RESET,
+    _SGLDiffusionLogger,
+    get_is_main_process,
+    init_logger,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass
+class MemorySnapshot:
+    allocated_mb: float  # current allocated memory
+    reserved_mb: float  # current reserved memory (actual VRAM)
+    peak_allocated_mb: float  # peak allocated since last reset
+    peak_reserved_mb: float  # peak reserved since last reset
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "allocated_mb": round(self.allocated_mb, 2),
+            "reserved_mb": round(self.reserved_mb, 2),
+            "peak_allocated_mb": round(self.peak_allocated_mb, 2),
+            "peak_reserved_mb": round(self.peak_reserved_mb, 2),
+        }
+
+
+@dataclasses.dataclass
+class RequestMetrics:
+    """Performance metrics for a single request, including timings and memory snapshots."""
+
+    def __init__(self, request_id: str):
+        self.request_id = request_id
+        self.stages: Dict[str, float] = {}
+        self.steps: list[float] = []
+        self.total_duration_ms: float = 0.0
+        self.suppress_stage_breakdown: bool = False
+        # memory tracking: {checkpoint_name: MemorySnapshot}
+        self.memory_snapshots: Dict[str, MemorySnapshot] = {}
+
+    @property
+    def total_duration_s(self) -> float:
+        return self.total_duration_ms / 1000.0
+
+    def record_stage(self, stage_name: str, duration_s: float):
+        """Records the duration of a pipeline stage"""
+        if self.suppress_stage_breakdown:
+            return
+        self.stages[stage_name] = duration_s * 1000  # Store as milliseconds
+
+    def record_step(self, duration_s: float):
+        """Records the duration of a denoising step in execution order."""
+        if self.suppress_stage_breakdown:
+            return
+        self.steps.append(duration_s * 1000)
+
+    def record_memory_snapshot(self, checkpoint_name: str, snapshot: MemorySnapshot):
+        if self.suppress_stage_breakdown:
+            return
+        self.memory_snapshots[checkpoint_name] = snapshot
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Serializes the metrics data to a dictionary."""
+        return {
+            "request_id": self.request_id,
+            "stages": self.stages,
+            "steps": self.steps,
+            "total_duration_ms": self.total_duration_ms,
+            "memory_snapshots": {
+                name: snapshot.to_dict()
+                for name, snapshot in self.memory_snapshots.items()
+            },
+        }
+
+
+def get_diffusion_perf_log_dir() -> str:
+    """
+    Determines the directory for performance logs.
+    """
+    log_dir = os.environ.get("SGLANG_PERF_LOG_DIR")
+    if log_dir:
+        return os.path.abspath(log_dir)
+    if log_dir is None:
+        sglang_path = Path(sglang.__file__).resolve()
+        target_path = (sglang_path.parent / "../../.cache/logs").resolve()
+        return str(target_path)
+    return ""
+
+
+@lru_cache(maxsize=1)
+def get_git_commit_hash() -> str:
+    try:
+        commit_hash = os.environ.get("SGLANG_GIT_COMMIT")
+        if not commit_hash:
+            commit_hash = (
+                subprocess.check_output(
+                    ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
+                )
+                .strip()
+                .decode("utf-8")
+            )
+        _CACHED_COMMIT_HASH = commit_hash
+        return commit_hash
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        _CACHED_COMMIT_HASH = "N/A"
+        return "N/A"
+
+
+def capture_memory_snapshot() -> MemorySnapshot:
+    if not torch.get_device_module().is_available():
+        return MemorySnapshot(
+            allocated_mb=0.0,
+            reserved_mb=0.0,
+            peak_allocated_mb=0.0,
+            peak_reserved_mb=0.0,
+        )
+
+    allocated = torch.get_device_module().memory_allocated()
+    reserved = torch.get_device_module().memory_reserved()
+    peak_allocated = torch.get_device_module().max_memory_allocated()
+    peak_reserved = torch.get_device_module().max_memory_reserved()
+
+    return MemorySnapshot(
+        allocated_mb=allocated / (1024**2),
+        reserved_mb=reserved / (1024**2),
+        peak_allocated_mb=peak_allocated / (1024**2),
+        peak_reserved_mb=peak_reserved / (1024**2),
+    )
+
+
+@dataclasses.dataclass
+class RequestPerfRecord:
+    request_id: str
+
+    timestamp: str
+    commit_hash: str
+    tag: str
+
+    stages: list[dict]
+    steps: list[float]
+    total_duration_ms: float
+    memory_snapshots: dict[str, dict] = dataclasses.field(default_factory=dict)
+
+    def __init__(
+        self,
+        request_id,
+        commit_hash,
+        tag,
+        stages,
+        steps,
+        total_duration_ms,
+        memory_snapshots=None,
+        timestamp=None,
+    ):
+        self.request_id = request_id
+        if timestamp is not None:
+            self.timestamp = timestamp
+        else:
+            self.timestamp = datetime.now(UTC).isoformat()
+
+        self.commit_hash = commit_hash
+        self.tag = tag
+        self.stages = stages
+        self.steps = steps
+        self.total_duration_ms = total_duration_ms
+        self.memory_snapshots = memory_snapshots or {}
+
+
+class StageProfiler:
+    """
+    A unified context manager, records performance metrics (usually of a single Stage or a step) into a provided RequestMetrics object (usually from a Req).
+    """
+
+    def __init__(
+        self,
+        stage_name: str,
+        logger: _SGLDiffusionLogger,
+        metrics: Optional["RequestMetrics"],
+        log_stage_start_end: bool = False,
+        perf_dump_path_provided: bool = False,
+        capture_memory: bool = False,
+        record_as_step: bool = False,
+    ):
+        self.stage_name = stage_name
+        self.metrics = metrics
+        self.logger = logger
+        self.start_time = 0.0
+        self.log_timing = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING
+        self.log_stage_start_end = log_stage_start_end
+        self.capture_memory = capture_memory
+        self.record_as_step = record_as_step
+
+    def _should_record_as_step(self) -> bool:
+        return self.record_as_step or self.stage_name.startswith("denoising_step_")
+
+    def __enter__(self):
+        if self.log_stage_start_end:
+            msg = f"[{self.stage_name}] started..."
+            if self.logger.isEnabledFor(logging.DEBUG):
+                # This debug-only memory log runs at every stage boundary in CI.
+                # Keep it observational; cache cleanup is handled at explicit
+                # failure and component-release points.
+                available_memory = current_platform.get_available_gpu_memory(
+                    empty_cache=False
+                )
+                msg += f" ({round(available_memory, 2)} GB left)"
+            self.logger.info(msg)
+
+        if (self.log_timing and self.metrics) or self.log_stage_start_end:
+            if (
+                os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
+                and self._should_record_as_step()
+                and torch.get_device_module().is_available()
+            ):
+                torch.get_device_module().synchronize()
+            self.start_time = time.perf_counter()
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if not ((self.log_timing and self.metrics) or self.log_stage_start_end):
+            return False
+
+        if (
+            os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
+            and self._should_record_as_step()
+            and torch.get_device_module().is_available()
+        ):
+            torch.get_device_module().synchronize()
+        execution_time_s = time.perf_counter() - self.start_time
+
+        if exc_type:
+            self.logger.error(
+                "[%s] Error during execution after %.4f ms: %s",
+                self.stage_name,
+                execution_time_s * 1000,
+                exc_val,
+                exc_info=True,
+            )
+            return False
+
+        if self.log_stage_start_end:
+            self.logger.info(
+                f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds",
+            )
+
+        if self.log_timing and self.metrics:
+            if self._should_record_as_step():
+                self.metrics.record_step(execution_time_s)
+            else:
+                self.metrics.record_stage(self.stage_name, execution_time_s)
+
+            # capture memory snapshot after stage if requested
+            if self.capture_memory and torch.get_device_module().is_available():
+                snapshot = capture_memory_snapshot()
+                self.metrics.record_memory_snapshot(
+                    f"after_{self.stage_name}", snapshot
+                )
+
+        return False
+
+
+class PerformanceLogger:
+    """
+    A global utility class for logging performance metrics for all request, categorized by request-id.
+
+    Serves both as a runtime logger (stream to file) and a dump utility.
+
+    Notice that RequestMetrics stores the performance metrics of a single request
+    """
+
+    @classmethod
+    def dump_benchmark_report(
+        cls,
+        file_path: str,
+        metrics: "RequestMetrics",
+        meta: Optional[Dict[str, Any]] = None,
+        tag: str = "benchmark_dump",
+    ):
+        """
+        Static method to dump a standardized benchmark report to a file.
+        Eliminates duplicate logic in CLI/Client code.
+        """
+        formatted_steps = [
+            {"name": name, "duration_ms": duration_ms}
+            for name, duration_ms in metrics.stages.items()
+        ]
+
+        denoise_steps_ms = [
+            {"step": idx, "duration_ms": duration_ms}
+            for idx, duration_ms in enumerate(metrics.steps)
+        ]
+
+        memory_checkpoints = {
+            name: snapshot.to_dict()
+            for name, snapshot in metrics.memory_snapshots.items()
+        }
+
+        report = {
+            "timestamp": datetime.now(UTC).isoformat(),
+            "request_id": metrics.request_id,
+            "commit_hash": get_git_commit_hash(),
+            "tag": tag,
+            "total_duration_ms": metrics.total_duration_ms,
+            "steps": formatted_steps,
+            "denoise_steps_ms": denoise_steps_ms,
+            "memory_checkpoints": memory_checkpoints,
+            "meta": meta or {},
+        }
+
+        try:
+            abs_path = os.path.abspath(file_path)
+            os.makedirs(os.path.dirname(abs_path), exist_ok=True)
+            with open(abs_path, "w", encoding="utf-8") as f:
+                json.dump(report, f, indent=2)
+            logger.info(f"Metrics dumped to: {CYAN}{abs_path}{RESET}")
+        except IOError as e:
+            logger.error(f"Failed to dump metrics to {abs_path}: {e}")
+
+    @classmethod
+    def log_request_summary(
+        cls,
+        metrics: "RequestMetrics",
+        tag: str = "total_inference_time",
+    ):
+        """logs the stage metrics and total duration for a completed request
+        to the performance_log file.
+
+        Note that this accords to the time spent internally in server, postprocess is not included
+        """
+        formatted_stages = [
+            {"name": name, "execution_time_ms": duration_ms}
+            for name, duration_ms in metrics.stages.items()
+        ]
+
+        memory_checkpoints = {
+            name: snapshot.to_dict()
+            for name, snapshot in metrics.memory_snapshots.items()
+        }
+
+        record = RequestPerfRecord(
+            metrics.request_id,
+            commit_hash=get_git_commit_hash(),
+            tag="pipeline_stage_metrics",
+            stages=formatted_stages,
+            steps=metrics.steps,
+            total_duration_ms=metrics.total_duration_ms,
+            memory_snapshots=memory_checkpoints,
+        )
+
+        try:
+            if get_is_main_process():
+                log_dir = get_diffusion_perf_log_dir()
+                if not os.path.exists(log_dir):
+                    os.makedirs(log_dir, exist_ok=True)
+
+                log_file = os.path.join(log_dir, "performance.log")
+
+                with open(log_file, "a", encoding="utf-8") as f:
+                    f.write(json.dumps(dataclasses.asdict(record)) + "\n")
+
+        except (OSError, PermissionError) as e:
+            print(f"WARNING: Failed to log performance record: {e}", file=sys.stderr)

python/sglang/multimodal_gen/runtime/utils/precision.py

@@ -0,0 +1,119 @@
+from contextlib import contextmanager
+from typing import Iterator, Optional, Union
+
+import torch
+
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+
+def precision_to_dtype(precision: str, field_name: str = "precision") -> torch.dtype:
+    try:
+        return PRECISION_TO_TYPE[precision]
+    except KeyError as exc:
+        raise ValueError(
+            f"Unsupported {field_name}={precision!r}; "
+            f"expected one of {sorted(PRECISION_TO_TYPE)}"
+        ) from exc
+
+
+def resolve_precision(
+    server_args,
+    component_or_precision_attr: str,
+    *,
+    precision_attr: Optional[str] = None,
+    field_name: Optional[str] = None,
+) -> torch.dtype:
+    precision_attr = precision_attr or component_or_precision_attr
+    precision = getattr(server_args.pipeline_config, precision_attr)
+    return precision_to_dtype(precision, field_name or precision_attr)
+
+
+def resolve_component_precision(server_args, module_name: str) -> Optional[torch.dtype]:
+    pipeline_config = getattr(server_args, "pipeline_config", None)
+    if pipeline_config is None:
+        return None
+
+    if module_name in ("audio_vae", "vocoder"):
+        precision_attr = "audio_vae_precision"
+    elif module_name in ("vae", "video_vae"):
+        precision_attr = "vae_precision"
+    elif module_name in (
+        "transformer",
+        "transformer_2",
+        "audio_dit",
+        "video_dit",
+        "connectors",
+        "dual_tower_bridge",
+    ):
+        precision_attr = "dit_precision"
+    elif module_name == "image_encoder":
+        precision_attr = "image_encoder_precision"
+    elif module_name == "text_encoder" or module_name.startswith("text_encoder_"):
+        precisions = getattr(pipeline_config, "text_encoder_precisions", None)
+        if not precisions:
+            return None
+        suffix = module_name.removeprefix("text_encoder")
+        index = 0 if suffix == "" else int(suffix.removeprefix("_")) - 1
+        if index < 0 or index >= len(precisions):
+            raise ValueError(
+                f"No configured precision for {module_name!r}; "
+                f"text_encoder_precisions has {len(precisions)} entries"
+            )
+        precision = precisions[index]
+        return precision_to_dtype(precision, f"text_encoder_precisions[{index}]")
+    else:
+        return None
+
+    if not hasattr(pipeline_config, precision_attr):
+        return None
+    return resolve_precision(server_args, precision_attr)
+
+
+def autocast_enabled(dtype: torch.dtype, disable_autocast: bool) -> bool:
+    return dtype != torch.float32 and not disable_autocast
+
+
+def get_module_dtype(module, default: torch.dtype = torch.float32) -> torch.dtype:
+    try:
+        return next(module.parameters()).dtype
+    except (AttributeError, StopIteration):
+        dtype = getattr(module, "dtype", None)
+        return dtype if isinstance(dtype, torch.dtype) else default
+
+
+def align_tensor_to_module_dtype(
+    tensor: torch.Tensor,
+    module,
+    *,
+    device: Optional[Union[torch.device, str]] = None,
+    default_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    dtype = get_module_dtype(module, default=default_dtype)
+    if device is None:
+        try:
+            device = next(module.parameters()).device
+        except (AttributeError, StopIteration):
+            device = tensor.device
+    if not tensor.is_floating_point():
+        return tensor.to(device=device)
+    return tensor.to(device=device, dtype=dtype)
+
+
+@contextmanager
+def temporary_module_dtype(
+    module,
+    dtype: torch.dtype,
+    *,
+    enabled: bool = True,
+    restore_dtype: Optional[torch.dtype] = None,
+) -> Iterator:
+    if not enabled:
+        yield module
+        return
+
+    original_dtype = restore_dtype or get_module_dtype(module)
+    module = module.to(dtype=dtype)
+    try:
+        yield module
+    finally:
+        module.to(dtype=original_dtype)

python/sglang/multimodal_gen/runtime/utils/profiler.py

@@ -0,0 +1,194 @@
+import gzip
+import os
+
+import torch
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger
+from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
+
+if current_platform.is_npu():
+    import torch_npu
+
+    patches = [
+        ["profiler.profile", torch_npu.profiler.profile],
+        ["profiler.schedule", torch_npu.profiler.schedule],
+    ]
+    apply_torch_npu_patches(torch_npu, patches)
+
+logger = init_logger(__name__)
+
+
+def _resolve_profiler_log_dir(log_dir: str | None) -> str:
+    if log_dir is not None:
+        return log_dir
+
+    diffusion_profiler_dir = os.getenv("SGLANG_DIFFUSION_TORCH_PROFILER_DIR")
+    if diffusion_profiler_dir:
+        return diffusion_profiler_dir
+
+    return os.getenv("SGLANG_TORCH_PROFILER_DIR", "./logs")
+
+
+class SGLDiffusionProfiler:
+    """
+    A wrapper around torch.profiler to simplify usage in pipelines.
+    Supports both full profiling and scheduled profiling.
+
+
+    1. if profile_all_stages is on: profile all stages, including all denoising steps
+    2. otherwise, if num_profiled_timesteps is specified: profile {num_profiled_timesteps} denoising steps. profile all steps if num_profiled_timesteps==-1
+    """
+
+    _instance = None
+
+    def __init__(
+        self,
+        request_id: str | None = None,
+        rank: int = 0,
+        full_profile: bool = False,
+        num_steps: int | None = None,
+        num_inference_steps: int | None = None,
+        log_dir: str | None = None,
+    ):
+        self.request_id = request_id or "profile_trace"
+        self.rank = rank
+        self.full_profile = full_profile
+
+        self.log_dir = _resolve_profiler_log_dir(log_dir)
+
+        try:
+            os.makedirs(self.log_dir, exist_ok=True)
+        except OSError:
+            pass
+
+        activities = [torch.profiler.ProfilerActivity.CPU]
+        if torch.cuda.is_available() or (
+            hasattr(torch, "musa") and torch.musa.is_available()
+        ):
+            activities.append(torch.profiler.ProfilerActivity.CUDA)
+        if current_platform.is_npu():
+            activities.append(torch_npu.profiler.ProfilerActivity.NPU)
+
+        if hasattr(torch, "xpu") and torch.xpu.is_available():
+            activities.append(torch.profiler.ProfilerActivity.XPU)
+
+        common_torch_profiler_args = dict(
+            activities=activities,
+            record_shapes=True,
+            with_stack=True,
+            on_trace_ready=(
+                None
+                if not current_platform.is_npu()
+                else torch_npu.profiler.tensorboard_trace_handler(self.log_dir)
+            ),
+        )
+        if self.full_profile:
+            # profile all stages
+            self.profiler = torch.profiler.profile(**common_torch_profiler_args)
+            self.profile_mode_id = "full stages"
+        else:
+            # profile denoising stage only
+            warmup = 1
+            num_actual_steps = num_inference_steps if num_steps == -1 else num_steps
+            self.num_active_steps = num_actual_steps + warmup
+            self.profiler = torch.profiler.profile(
+                **common_torch_profiler_args,
+                schedule=torch.profiler.schedule(
+                    skip_first=0,
+                    wait=0,
+                    warmup=warmup,
+                    active=self.num_active_steps,
+                    repeat=1,
+                ),
+            )
+            self.profile_mode_id = f"{num_actual_steps} steps"
+
+        logger.info(f"Profiling request: {request_id} for {self.profile_mode_id}...")
+
+        self.has_stopped = False
+
+        SGLDiffusionProfiler._instance = self
+        self.start()
+
+    def start(self):
+        logger.info("Starting Profiler...")
+        self.profiler.start()
+
+    def _step(self):
+        self.profiler.step()
+
+    def step_stage(self):
+        if self.full_profile:
+            self._step()
+
+    def step_denoising_step(self):
+        if not self.full_profile:
+            if self.num_active_steps >= 0:
+                self._step()
+                self.num_active_steps -= 1
+            else:
+                # early exit when enough steps are captured, to reduce the trace file size
+                self.stop(dump_rank=0)
+
+    @classmethod
+    def get_instance(cls) -> "SGLDiffusionProfiler":
+        return cls._instance
+
+    def stop(self, export_trace: bool = True, dump_rank: int | None = None):
+        if self.has_stopped:
+            return
+        self.has_stopped = True
+        logger.info("Stopping Profiler...")
+        if torch.cuda.is_available() or (
+            hasattr(torch, "musa") and torch.musa.is_available()
+        ):
+            torch.cuda.synchronize()
+        if current_platform.is_npu():
+            torch.npu.synchronize()
+            export_trace = False  # set to false because our internal torch_npu.profiler will generate trace file
+        self.profiler.stop()
+
+        if export_trace:
+            if dump_rank is not None and dump_rank != self.rank:
+                pass
+            else:
+                self._export_trace()
+
+        SGLDiffusionProfiler._instance = None
+
+    def _export_trace(self):
+
+        try:
+            os.makedirs(self.log_dir, exist_ok=True)
+            sanitized_profile_mode_id = self.profile_mode_id.replace(" ", "_")
+            trace_path = os.path.abspath(
+                os.path.join(
+                    self.log_dir,
+                    f"{self.request_id}-{sanitized_profile_mode_id}-global-rank{self.rank}.trace.json.gz",
+                )
+            )
+            self.profiler.export_chrome_trace(trace_path)
+
+            if self._check_trace_integrity(trace_path):
+                logger.info(f"Saved profiler traces to: {CYAN}{trace_path}{RESET}")
+            else:
+                logger.warning(f"Trace file may be corrupted: {trace_path}")
+        except Exception as e:
+            logger.error(f"Failed to save trace: {e}")
+
+    def _check_trace_integrity(self, trace_path: str) -> bool:
+        try:
+            if not os.path.exists(trace_path) or os.path.getsize(trace_path) == 0:
+                return False
+
+            with gzip.open(trace_path, "rb") as f:
+                content = f.read()
+                if content.count(b"\x1f\x8b") > 1:
+                    logger.warning("Multiple gzip headers detected")
+                    return False
+
+            return True
+        except Exception as e:
+            logger.warning(f"Trace file integrity check failed: {e}")
+            return False

python/sglang/multimodal_gen/runtime/utils/quantization_utils.py

@@ -0,0 +1,567 @@
+import glob
+import json
+import os
+import re
+import struct
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from safetensors import safe_open
+
+from sglang.multimodal_gen.runtime.layers.quantization import (
+    QuantizationConfig,
+    get_quantization_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def normalize_flat_modelopt_quant_config(
+    quant_cfg: dict[str, Any] | None,
+) -> dict[str, Any] | None:
+    """Fill required diffusers fields for flat ModelOpt component configs."""
+    if not isinstance(quant_cfg, dict) or quant_cfg.get("quant_method") != "modelopt":
+        return quant_cfg
+
+    quant_algo = str(
+        quant_cfg.get("quant_algo")
+        or quant_cfg.get("quantization", {}).get("quant_algo")
+        or ""
+    ).upper()
+    if not quant_algo:
+        return quant_cfg
+
+    normalized = dict(quant_cfg)
+    normalized.setdefault("quant_type", quant_algo)
+    return normalized
+
+
+def _infer_nvfp4_group_size_from_tensors(weight, scale) -> Optional[int]:
+    """Infer NVFP4 group_size from serialized weight/scale tensor shapes."""
+    return _infer_nvfp4_group_size_from_shapes(
+        getattr(weight, "shape", ()),
+        getattr(scale, "shape", ()),
+    )
+
+
+def _infer_nvfp4_group_size_from_shapes(weight_shape, scale_shape) -> Optional[int]:
+    weight_shape = tuple(weight_shape or ())
+    scale_shape = tuple(scale_shape or ())
+    if len(weight_shape) < 2:
+        return None
+
+    input_size = int(weight_shape[1]) * 2
+    if input_size <= 0:
+        return None
+
+    candidate_num_groups: list[int] = []
+    if len(scale_shape) >= 2:
+        candidate_num_groups.append(int(scale_shape[-1]))
+    elif len(scale_shape) == 1:
+        scale_len = int(scale_shape[0])
+        if scale_len == int(weight_shape[0]):
+            candidate_num_groups.append(1)
+        candidate_num_groups.append(scale_len)
+    else:
+        candidate_num_groups.append(1)
+
+    for num_groups in candidate_num_groups:
+        if num_groups <= 0:
+            continue
+        if input_size % num_groups == 0:
+            return input_size // num_groups
+
+    return None
+
+
+def _read_safetensors_tensor_metadata(file_path: str) -> dict[str, dict[str, Any]]:
+    with open(file_path, "rb") as f:
+        header_len = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(header_len))
+    header.pop("__metadata__", None)
+    return header
+
+
+def _is_nvfp4_tensor_family(
+    module_name: str,
+    tensor_metadata: dict[str, dict[str, Any]],
+) -> bool:
+    weight_metadata = tensor_metadata.get(f"{module_name}.weight")
+    scale_metadata = tensor_metadata.get(f"{module_name}.weight_scale")
+    if weight_metadata is None or scale_metadata is None:
+        return False
+
+    weight_dtype = str(weight_metadata.get("dtype", "")).upper()
+    scale_dtype = str(scale_metadata.get("dtype", "")).upper()
+    scale_shape = scale_metadata.get("shape", [])
+    return weight_dtype == "U8" and "F8_E4M3" in scale_dtype and len(scale_shape) >= 2
+
+
+def _resolve_quant_method_name(quant_cfg: dict) -> str:
+    quant_cfg = normalize_flat_modelopt_quant_config(quant_cfg) or quant_cfg
+    quant_method = quant_cfg.get("quant_method")
+    if quant_method == "bitsandbytes":
+        return "bitsandbytes"
+    if quant_method != "modelopt":
+        return quant_method
+
+    quant_algo = (
+        quant_cfg.get("quant_algo")
+        or quant_cfg.get("quantization", {}).get("quant_algo")
+        or ""
+    ).upper()
+    if quant_algo == "MIXED_PRECISION":
+        raise ValueError(
+            "ModelOpt mixed precision is not supported by the current SGLang diffusion runtime."
+        )
+    if "FP8" in quant_algo:
+        return "modelopt_fp8"
+    if "FP4" in quant_algo or "NVFP4" in quant_algo:
+        return "modelopt_fp4"
+    raise ValueError(f"Unsupported ModelOpt quant_algo for diffusion: {quant_algo}")
+
+
+def _load_quant_cls(quant_cfg: dict):
+    quant_method = _resolve_quant_method_name(quant_cfg)
+    if not quant_method:
+        raise ValueError("Missing quant_method in quantization config.")
+    return get_quantization_config(quant_method)
+
+
+def find_quant_modelslim_config(model_config, component_model_path):
+    # Try exact name first, then glob for variant filenames (e.g. after repack)
+    quant_config_file = Path(component_model_path, "quant_model_description.json")
+    if not quant_config_file.is_file():
+        candidates = sorted(
+            Path(component_model_path).glob("quant_model_description*.json")
+        )
+        quant_config_file = candidates[0] if candidates else None
+
+    quant_cfg = None
+    if quant_config_file is not None and Path(quant_config_file).is_file():
+        with open(quant_config_file) as f:
+            quant_cfg = json.load(f)
+        # This field is required for flagless model loading but is not present in
+        # modelslim model description, so we're adding it here manually.
+        quant_cfg["quant_method"] = "modelslim"
+
+    return quant_cfg
+
+
+def replace_prefix(key: str, prefix_mapping: dict[str, str]) -> str:
+    for prefix, new_prefix in prefix_mapping.items():
+        if key.startswith(prefix):
+            key = key.replace(prefix, new_prefix, 1)
+    return key
+
+
+def get_quant_config(
+    model_config,
+    component_model_path: str,
+    packed_modules_mapping: Dict[str, List[str]] = {},
+    reverse_param_names_mapping: Dict[str, List[str]] = {},
+    remap_prefix: Dict[str, str] | None = None,
+    quant_ignore_remap: Optional[Dict[str, str]] = None,
+) -> QuantizationConfig:
+    quant_cfg = find_quant_modelslim_config(model_config, component_model_path)
+    if quant_cfg is not None:
+        quant_cls = _load_quant_cls(quant_cfg)
+        return quant_cls.from_config(quant_cfg, reverse_param_names_mapping)
+
+    if "quantization_config" not in model_config:
+        return None
+
+    hf_quant_config = normalize_flat_modelopt_quant_config(
+        model_config["quantization_config"]
+    )
+    if hf_quant_config is not None and not isinstance(hf_quant_config, dict):
+        hf_quant_config = hf_quant_config.to_dict()
+    quant_cls = _load_quant_cls(hf_quant_config)
+
+    # GGUF doesn't have config file
+    if hf_quant_config["quant_method"] == "gguf":
+        return quant_cls.from_config({})
+
+    # some vision model may keep quantization_config in their text_config
+    hf_text_config = getattr(model_config, "text_config", None)
+    if hf_quant_config is None and hf_text_config is not None:
+        hf_quant_config = getattr(hf_text_config, "quantization_config", None)
+    if hf_quant_config is None:
+        # compressed-tensors uses a compressions_config
+        hf_quant_config = getattr(model_config, "compression_config", None)
+    if hf_quant_config is not None:
+        hf_quant_config["packed_modules_mapping"] = packed_modules_mapping
+        is_modelopt_fp8 = (
+            hf_quant_config.get("quant_method") == "modelopt"
+            and "FP8" in str(hf_quant_config.get("quant_algo", "")).upper()
+        )
+        extra_kwargs = (
+            {"ignore_remap": quant_ignore_remap}
+            if quant_ignore_remap and is_modelopt_fp8
+            else {}
+        )
+        return quant_cls.from_config(hf_quant_config, **extra_kwargs)
+
+    model_name_or_path = model_config["model_path"]
+    hf_folder = model_name_or_path
+
+    possible_config_filenames = quant_cls.get_config_filenames()
+
+    # If the quantization config is not found, use the default config.
+    if not possible_config_filenames:
+        return quant_cls()
+
+    config_files = glob.glob(os.path.join(hf_folder, "*.json"))
+
+    quant_config_files = [
+        f for f in config_files if any(f.endswith(x) for x in possible_config_filenames)
+    ]
+    if len(quant_config_files) == 0:
+        raise ValueError(
+            f"Cannot find the config file for {model_config['quantization_config']['quant_method']}"
+        )
+    if len(quant_config_files) > 1:
+        raise ValueError(
+            f"Found multiple config files for {model_config['quantization_config']['quant_method']}: "
+            f"{quant_config_files}"
+        )
+
+    quant_config_file = quant_config_files[0]
+    with open(quant_config_file) as f:
+        config = json.load(f)
+        if remap_prefix is not None and "quantization" in config:
+            exclude_modules = [
+                replace_prefix(key, remap_prefix)
+                for key in config["quantization"]["exclude_modules"]
+            ]
+            config["quantization"]["exclude_modules"] = exclude_modules
+        config["packed_modules_mapping"] = packed_modules_mapping
+        return quant_cls.from_config(config)
+
+
+def handle_fp8_metadata_format(quant_config_dict):
+    layers = quant_config_dict.get("layers", {})
+    if any(
+        isinstance(v, dict) and "float8" in v.get("format", "") for v in layers.values()
+    ):
+        quant_config_dict["quant_method"] = "fp8"
+        quant_config_dict["activation_scheme"] = "dynamic"
+    return quant_config_dict
+
+
+def get_quant_config_from_safetensors_metadata(
+    file_path: str,
+) -> Optional[QuantizationConfig]:
+    """Extract quantization config from a safetensors file's metadata header.
+    Returns None if no recognizable quantization metadata is found.
+    """
+    metadata = get_metadata_from_safetensors_file(file_path)
+    if not metadata:
+        return None
+
+    quant_config_str = metadata.get("_quantization_metadata")
+    quant_config_dict = None
+    if quant_config_str:
+        try:
+            quant_config_dict = json.loads(quant_config_str)
+        except Exception:
+            quant_config_dict = None
+
+    if quant_config_dict is None:
+        quant_config_str = metadata.get("quantization_config")
+        if not quant_config_str:
+            return None
+        try:
+            quant_config_dict = json.loads(quant_config_str)
+        except Exception:
+            return None
+
+    if not quant_config_dict:
+        return None
+
+    # handle diffusers fp8 safetensors metadata format
+    if (
+        "quant_method" not in quant_config_dict
+        and "format_version" in quant_config_dict
+        and "layers" in quant_config_dict
+    ):
+        quant_config_dict = handle_fp8_metadata_format(quant_config_dict)
+
+    quant_method = quant_config_dict.get("quant_method")
+    if not quant_method:
+        return None
+
+    try:
+        quant_cls = _load_quant_cls(quant_config_dict)
+        config = quant_cls.from_config(quant_config_dict)
+        logger.debug(f"Get quantization config from safetensors file: {file_path}")
+        return config
+    except Exception as _e:
+        return None
+
+
+def get_metadata_from_safetensors_file(file_path: str):
+    try:
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            metadata = f.metadata()
+            return metadata
+    except Exception as e:
+        logger.warning(e)
+
+
+def _canonicalize_modulation_exclude(module_name: str) -> str:
+    """Map a serialized modulation weight's parent to the runtime linear prefix.
+
+    Qwen-Image wraps the modulation projection in ``nn.Sequential(SiLU, Linear)``,
+    so its weights serialize as ``...img_mod.1.weight`` while the runtime
+    ReplicatedLinear advertises ``...img_mod`` as its quant/exclusion prefix.
+    Strip the trailing Sequential index so a safetensors-inferred BF16 exclude
+    entry actually matches the linear (mirrors the ModelOpt FP8 converter, which
+    canonicalizes ``.img_mod.1``/``.txt_mod.1`` to ``.img_mod``/``.txt_mod``).
+    No-op for any other module name.
+    """
+    if module_name.endswith((".img_mod.1", ".txt_mod.1")):
+        return module_name.removesuffix(".1")
+    return module_name
+
+
+def _build_nvfp4_config_from_safetensors_files(
+    file_paths: list[str],
+    param_names_mapping_dict: Optional[dict] = None,
+    reverse_param_names_mapping_dict: Optional[dict] = None,
+    fallback_group_size: Optional[int] = None,
+) -> Optional[QuantizationConfig]:
+    """Build a single NVFP4 config by aggregating metadata across multiple files.
+
+    Some checkpoints split BF16 fallback layers and NVFP4 layers across multiple
+    safetensors. Building the config from only the first matching file can
+    incorrectly exclude layers that are quantized in a later shard.
+    """
+    group_size = None
+    quantized_bfl_modules: set[str] = set()
+    non_quantized_bfl_modules: set[str] = set()
+    files_with_nvfp4_signal: list[str] = []
+    checkpoint_uses_packed_qkv = False
+    checkpoint_uses_comfy_quant = False
+    packed_qkv_pattern = re.compile(
+        r"^(double_blocks\.\d+\.(img|txt)_attn\.qkv|single_blocks\.\d+\.linear1)\."
+    )
+
+    for file_path in file_paths:
+        metadata = get_metadata_from_safetensors_file(file_path)
+        quant_config_dict = None
+        metadata_signals_nvfp4 = False
+        if metadata:
+            quant_config_str = metadata.get("_quantization_metadata")
+            if quant_config_str:
+                try:
+                    quant_config_dict = json.loads(quant_config_str)
+                except json.JSONDecodeError:
+                    quant_config_dict = None
+                else:
+                    quant_algo = str(quant_config_dict.get("quant_algo", "")).upper()
+                    quant_type = str(quant_config_dict.get("quant_type", "")).upper()
+                    metadata_signals_nvfp4 = (
+                        "NVFP4" in quant_algo
+                        or "FP4" in quant_algo
+                        or "NVFP4" in quant_type
+                    )
+
+        file_quantized_modules: set[str] = set()
+        if (
+            quant_config_dict is not None
+            and "format_version" in quant_config_dict
+            and "layers" in quant_config_dict
+        ):
+            layers = quant_config_dict.get("layers", {})
+            file_quantized_modules.update(
+                layer_name
+                for layer_name, layer_cfg in layers.items()
+                if isinstance(layer_cfg, dict) and layer_cfg.get("format") == "nvfp4"
+            )
+
+        tensor_metadata = _read_safetensors_tensor_metadata(file_path)
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            all_keys = set(f.keys())
+            if any(packed_qkv_pattern.match(k) for k in all_keys):
+                checkpoint_uses_packed_qkv = True
+            if any(k.endswith(".comfy_quant") for k in all_keys):
+                checkpoint_uses_comfy_quant = True
+
+            # Some ModelOpt NVFP4 exports only store a flat config.json plus
+            # per-file metadata without the diffusers `layers` section. Infer
+            # quantized modules directly from tensor families in that case.
+            # Mixed checkpoints may also contain FP8 fallback layers with scalar
+            # `.weight_scale`, so require packed uint8 weights and block scales.
+            file_quantized_modules.update(
+                key[: -len(".weight_scale")]
+                for key in all_keys
+                if key.endswith(".weight_scale")
+                and _is_nvfp4_tensor_family(
+                    key[: -len(".weight_scale")], tensor_metadata
+                )
+            )
+
+            if file_quantized_modules or metadata_signals_nvfp4:
+                files_with_nvfp4_signal.append(file_path)
+            quantized_bfl_modules.update(file_quantized_modules)
+
+            if group_size is None:
+                for layer_name in sorted(file_quantized_modules):
+                    weight_key = f"{layer_name}.weight"
+                    scale_key = f"{layer_name}.weight_scale"
+                    weight_metadata = tensor_metadata.get(weight_key)
+                    scale_metadata = tensor_metadata.get(scale_key)
+                    if weight_metadata is not None and scale_metadata is not None:
+                        group_size = _infer_nvfp4_group_size_from_shapes(
+                            weight_metadata.get("shape"),
+                            scale_metadata.get("shape"),
+                        )
+                        if group_size is not None:
+                            break
+
+            for k in sorted(all_keys):
+                if not k.endswith(".weight"):
+                    continue
+                module_name = k[: -len(".weight")]
+                if module_name not in file_quantized_modules:
+                    non_quantized_bfl_modules.add(module_name)
+
+    if not files_with_nvfp4_signal:
+        return None
+
+    if (
+        group_size is not None
+        and fallback_group_size is not None
+        and group_size != fallback_group_size
+    ):
+        logger.warning(
+            "NVFP4 group_size inferred from safetensors (%d) does not match config (%d); "
+            "preferring safetensors.",
+            group_size,
+            fallback_group_size,
+        )
+
+    if group_size is None and fallback_group_size is not None:
+        logger.info(
+            "Falling back to config-derived NVFP4 group_size=%d for %s",
+            fallback_group_size,
+            ", ".join(files_with_nvfp4_signal),
+        )
+        group_size = fallback_group_size
+
+    if group_size is None:
+        logger.warning(
+            "Could not infer group_size from NVFP4 safetensors: %s",
+            ", ".join(files_with_nvfp4_signal),
+        )
+        return None
+
+    exclude_bfl_modules = sorted(non_quantized_bfl_modules - quantized_bfl_modules)
+
+    exclude_modules = []
+    mapping_fn = None
+    reverse_mapping_fn = None
+    if param_names_mapping_dict or reverse_param_names_mapping_dict:
+        from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+
+        if param_names_mapping_dict:
+            mapping_fn = get_param_names_mapping(param_names_mapping_dict)
+        if reverse_param_names_mapping_dict:
+            reverse_mapping_fn = get_param_names_mapping(
+                reverse_param_names_mapping_dict
+            )
+
+    for module_bfl in exclude_bfl_modules:
+        raw_weight_name = f"{module_bfl}.weight"
+        if mapping_fn is not None:
+            mapped, _, _ = mapping_fn(raw_weight_name)
+            if mapped != raw_weight_name:
+                exclude_modules.append(
+                    mapped[: -len(".weight")] if mapped.endswith(".weight") else mapped
+                )
+                continue
+
+        if reverse_mapping_fn is not None:
+            reverse_mapped, _, _ = reverse_mapping_fn(raw_weight_name)
+            if reverse_mapped != raw_weight_name:
+                exclude_modules.append(
+                    reverse_mapped[: -len(".weight")]
+                    if reverse_mapped.endswith(".weight")
+                    else reverse_mapped
+                )
+                continue
+
+        exclude_modules.append(module_bfl)
+
+    exclude_modules = sorted(
+        {_canonicalize_modulation_exclude(m) for m in exclude_modules}
+    )
+
+    try:
+        quant_cls = get_quantization_config("modelopt_fp4")
+        checkpoint_uses_swizzled_scales = (
+            checkpoint_uses_packed_qkv or checkpoint_uses_comfy_quant
+        )
+        result = quant_cls.from_config(
+            {
+                "quant_algo": "NVFP4",
+                "group_size": group_size,
+                "ignore": exclude_modules,
+                "checkpoint_uses_packed_qkv": checkpoint_uses_packed_qkv,
+                # packed-QKV and Comfy NVFP4 checkpoints store serialized
+                # weights/scales in the FlashInfer/CUTLASS checkpoint layout
+                "checkpoint_weight_scale_layout": (
+                    "swizzled" if checkpoint_uses_swizzled_scales else "linear"
+                ),
+                "swap_weight_nibbles": checkpoint_uses_swizzled_scales,
+            }
+        )
+        logger.info(
+            "Built NVFP4 quant config from %d safetensors: group_size=%d, %d excluded modules, packed_qkv=%s, comfy_quant=%s, scale_layout=%s, swap_nibbles=%s",
+            len(files_with_nvfp4_signal),
+            group_size,
+            len(exclude_modules),
+            checkpoint_uses_packed_qkv,
+            checkpoint_uses_comfy_quant,
+            getattr(result, "checkpoint_weight_scale_layout", "linear"),
+            getattr(result, "swap_weight_nibbles", False),
+        )
+        return result
+    except Exception as e:
+        logger.warning(
+            "Failed to build NVFP4 config from %s: %s",
+            ", ".join(files_with_nvfp4_signal),
+            e,
+        )
+        return None
+
+
+def build_nvfp4_config_from_safetensors(
+    file_path: str,
+    param_names_mapping_dict: Optional[dict] = None,
+    reverse_param_names_mapping_dict: Optional[dict] = None,
+    fallback_group_size: Optional[int] = None,
+) -> Optional[QuantizationConfig]:
+    """Backward-compatible wrapper for a single safetensors file."""
+    return _build_nvfp4_config_from_safetensors_files(
+        [file_path],
+        param_names_mapping_dict,
+        reverse_param_names_mapping_dict,
+        fallback_group_size,
+    )
+
+
+def build_nvfp4_config_from_safetensors_list(
+    file_paths: list[str],
+    param_names_mapping_dict: Optional[dict] = None,
+    reverse_param_names_mapping_dict: Optional[dict] = None,
+    fallback_group_size: Optional[int] = None,
+) -> Optional[QuantizationConfig]:
+    return _build_nvfp4_config_from_safetensors_files(
+        file_paths,
+        param_names_mapping_dict,
+        reverse_param_names_mapping_dict,
+        fallback_group_size,
+    )

python/sglang/multimodal_gen/runtime/utils/realtime_video.py

@@ -0,0 +1,215 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import time
+import zlib
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
+        OutputBatch,
+        Req,
+    )
+
+logger = init_logger(__name__)
+
+RAW_RGB_CONTENT_TYPE = "application/x-raw-rgb"
+RAW_RGB_DELTA_GZIP_CONTENT_TYPE = "application/x-raw-rgb-delta-gzip"
+RAW_RGBA_DELTA_GZIP_CONTENT_TYPE = "application/x-raw-rgba-delta-gzip"
+WEBP_FRAME_CONTENT_TYPE = "image/webp"
+JPEG_FRAME_CONTENT_TYPE = "image/jpeg"
+RAW_RGB_CHANNELS = 3
+RAW_RGBA_CHANNELS = 4
+_RAW_RGB_DELTA_GZIP_LEVEL = 0
+
+
+def build_delta_gzip_raw_rgb_payload(
+    frames: list[bytes],
+    *,
+    reference_frame: bytes | None = None,
+) -> bytes:
+    if not frames:
+        return b""
+
+    frame_size = len(frames[0])
+    if reference_frame is not None and len(reference_frame) != frame_size:
+        raise ValueError("raw RGB delta gzip reference frame size mismatch")
+
+    previous = (
+        np.frombuffer(reference_frame, dtype=np.uint8)
+        if reference_frame is not None
+        else None
+    )
+    # keep gzip framing for lossless transport without spending realtime budget on compression
+    compressor = zlib.compressobj(
+        level=_RAW_RGB_DELTA_GZIP_LEVEL, method=zlib.DEFLATED, wbits=31
+    )
+    compressed_chunks = []
+    for frame in frames:
+        if len(frame) != frame_size:
+            raise ValueError("raw RGB delta gzip requires fixed-size frames")
+        current = np.frombuffer(frame, dtype=np.uint8)
+        if previous is None:
+            delta_frame = frame
+        else:
+            delta_frame = np.bitwise_xor(current, previous).tobytes()
+        compressed_chunks.append(compressor.compress(delta_frame))
+        previous = current
+
+    compressed_chunks.append(compressor.flush())
+    return b"".join(compressed_chunks)
+
+
+def restore_delta_gzip_raw_rgb_payload(
+    payload: bytes,
+    *,
+    bytes_per_frame: int,
+    num_frames: int,
+    reference_frame: bytes | None = None,
+) -> bytes:
+    if reference_frame is not None and len(reference_frame) != bytes_per_frame:
+        raise ValueError("delta gzip reference frame size mismatch")
+
+    delta_payload = zlib.decompress(payload, wbits=31)
+    expected_size = bytes_per_frame * num_frames
+    if len(delta_payload) != expected_size:
+        raise ValueError(
+            "delta gzip payload size mismatch: "
+            f"expected {expected_size}, got {len(delta_payload)}"
+        )
+
+    restored = bytearray(delta_payload)
+    previous = (
+        np.frombuffer(reference_frame, dtype=np.uint8)
+        if reference_frame is not None
+        else None
+    )
+    for frame_idx in range(num_frames):
+        offset = frame_idx * bytes_per_frame
+        current = np.frombuffer(
+            restored, dtype=np.uint8, count=bytes_per_frame, offset=offset
+        )
+        if previous is not None:
+            current ^= previous
+        previous = current
+    return bytes(restored)
+
+
+def build_raw_rgb_frame_batches(
+    output: Any,
+    req: Req,
+    output_batch: OutputBatch,
+    post_process_sample_fn: Callable[..., Any],
+) -> tuple[list[list[bytes]], dict[str, Any]]:
+    """post-process for realtime responses, returns only the batched frames and metadata"""
+    start = time.monotonic()
+    sample_to_frames_ms = 0.0
+    frames_to_bytes_ms = 0.0
+    raw_bytes = 0
+    num_frames = 0
+    frame_shape = None
+    frame_batches = []
+    if isinstance(output, torch.Tensor):
+        outputs = list(output)
+    else:
+        outputs = output if isinstance(output, Sequence) else [output]
+
+    for sample in outputs:
+        stage_start = time.monotonic()
+        if (
+            isinstance(sample, torch.Tensor)
+            and not req.enable_frame_interpolation
+            and not req.enable_upscaling
+        ):
+            frames = _tensor_sample_to_rgb24_array(sample)
+        else:
+            frames = post_process_sample_fn(
+                sample,
+                req.data_type,
+                req.fps,
+                False,
+                None,
+                audio_sample_rate=output_batch.audio_sample_rate,
+                output_compression=req.output_compression,
+                enable_frame_interpolation=req.enable_frame_interpolation,
+                frame_interpolation_exp=req.frame_interpolation_exp,
+                frame_interpolation_scale=req.frame_interpolation_scale,
+                frame_interpolation_model_path=req.frame_interpolation_model_path,
+                enable_upscaling=False,
+                upscaling_model_path=req.upscaling_model_path,
+                upscaling_scale=req.upscaling_scale,
+            )
+            if req.enable_upscaling and frames:
+                from sglang.multimodal_gen.runtime.postprocess import (
+                    batch_upscale_frames,
+                )
+
+                frames = batch_upscale_frames(
+                    frames,
+                    model_path=req.upscaling_model_path,
+                    scale=req.upscaling_scale,
+                )
+        sample_to_frames_ms += (time.monotonic() - stage_start) * 1000.0
+
+        stage_start = time.monotonic()
+
+        # numpy frames to RGB24 bytes
+        raw_frames = []
+        for frame in frames:
+            if frame.ndim == 2:
+                frame = frame[:, :, None]
+            if frame.shape[-1] == 1:
+                frame = np.repeat(frame, 3, axis=-1)
+            elif frame.shape[-1] > RAW_RGB_CHANNELS:
+                frame = frame[:, :, :RAW_RGB_CHANNELS]
+            frame = np.ascontiguousarray(frame)
+            frame_shape = tuple(int(dim) for dim in frame.shape)
+            frame_bytes = frame.tobytes()
+            raw_bytes += len(frame_bytes)
+            num_frames += 1
+            raw_frames.append(frame_bytes)
+        frames_to_bytes_ms += (time.monotonic() - stage_start) * 1000.0
+        frame_batches.append(raw_frames)
+
+    total_ms = (time.monotonic() - start) * 1000.0
+    logger.info(
+        "realtime raw RGB frame batch timing: request_id=%s "
+        "chunk_idx=%s sample_to_frames=%.2fms frames_to_bytes=%.2fms "
+        "total=%.2fms batches=%d frames=%d frame_shape=%s "
+        "raw_bytes=%d content_type=%s",
+        req.request_id,
+        req.block_idx,
+        sample_to_frames_ms,
+        frames_to_bytes_ms,
+        total_ms,
+        len(frame_batches),
+        num_frames,
+        frame_shape,
+        raw_bytes,
+        RAW_RGB_CONTENT_TYPE,
+    )
+    frame_metadata: dict[str, Any] = {}
+    if frame_shape is not None and len(frame_shape) == 3:
+        frame_height, frame_width, channels = frame_shape
+        frame_metadata = {
+            "format": "rgb24",
+            "width": frame_width,
+            "height": frame_height,
+            "channels": channels,
+            "bytes_per_frame": frame_width * frame_height * channels,
+        }
+    return frame_batches, frame_metadata
+
+
+def _tensor_sample_to_rgb24_array(sample: torch.Tensor) -> np.ndarray:
+    if sample.dim() == 3:
+        sample = sample.unsqueeze(1)
+    sample = (sample * 255).clamp(0, 255).to(torch.uint8)
+    return sample.permute(1, 2, 3, 0).contiguous().cpu().numpy()

python/sglang/multimodal_gen/runtime/utils/request_logger.py

@@ -0,0 +1,204 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+from typing import Any, Optional
+
+from sglang.srt.environ import envs
+from sglang.srt.utils.log_utils import create_log_targets, log_json
+from sglang.srt.utils.request_logger import (
+    _dataclass_to_string_truncated,
+    _transform_data_for_logging,
+)
+
+# Core generation knobs logged per record. Prompt text is logged separately,
+# gated by the level, so it is excluded here.
+_SAMPLING_CONFIG_FIELDS = (
+    "data_type",
+    "seed",
+    "num_inference_steps",
+    "guidance_scale",
+    "true_cfg_scale",
+    "width",
+    "height",
+    "num_frames",
+    "fps",
+    "num_outputs_per_prompt",
+)
+
+# Level 2 truncates prompt text; level 3 keeps it whole. Lower levels log no
+# prompt at all.
+_TRUNCATE_LENGTH = 2048
+_UNLIMITED = 1 << 30
+
+
+class DiffusionRequestLogger:
+    def __init__(
+        self,
+        log_requests: bool,
+        log_requests_level: int,
+        log_requests_format: str,
+        log_requests_target: Optional[list],
+    ):
+        self.log_requests = log_requests
+        self.log_requests_level = log_requests_level
+        self.log_requests_format = log_requests_format
+        self.log_requests_target = log_requests_target
+        self.targets = create_log_targets(
+            targets=log_requests_target, name_prefix=__name__
+        )
+        self.log_exceeded_ms = envs.SGLANG_LOG_REQUEST_EXCEEDED_MS.get()
+        self._max_length = (
+            _TRUNCATE_LENGTH if self.log_requests_level == 2 else _UNLIMITED
+        )
+
+    @classmethod
+    def from_server_args(cls, server_args: Any) -> "DiffusionRequestLogger":
+        """Build a logger from server args."""
+        return cls(
+            log_requests=server_args.log_requests,
+            log_requests_level=server_args.log_requests_level,
+            log_requests_format=server_args.log_requests_format,
+            log_requests_target=server_args.log_requests_target,
+        )
+
+    @staticmethod
+    def _request_id(req: Any) -> Optional[str]:
+        """The request's id, or ``None`` if absent."""
+        return getattr(req, "request_id", None)
+
+    def _config_view(self, req: Any, *, drop_seed: bool = False) -> dict:
+        """Sampling config + (level >= 2) prompt, gated by the log level.
+        Returns ``{}`` below level 1."""
+        sp = getattr(req, "sampling_params", None)
+        if sp is None or self.log_requests_level < 1:
+            return {}
+        cfg = {name: getattr(sp, name, None) for name in _SAMPLING_CONFIG_FIELDS}
+        if drop_seed:
+            cfg.pop("seed", None)
+        view: dict = {"sampling_params": cfg}
+        if self.log_requests_level >= 2:
+            view["prompt"] = getattr(sp, "prompt", None)
+            view["negative_prompt"] = getattr(sp, "negative_prompt", None)
+        return view
+
+    def _result_view(self, result: Any) -> dict:
+        """Result-side fields for a finished record: latency and error."""
+        e2e_latency = 0.0
+        metrics = getattr(result, "metrics", None) if result is not None else None
+        if metrics is not None:
+            e2e_latency = getattr(metrics, "total_duration_s", 0.0) or 0.0
+        return {
+            "meta_info": {"e2e_latency": e2e_latency},
+            "error": getattr(result, "error", None) if result is not None else None,
+        }
+
+    def _emit(self, msg: str) -> None:
+        for target in self.targets:
+            target.info(msg)
+
+    def _per_request_view(self, req: Any) -> dict:
+        """Per-output identity within a batch: ``request_id`` plus ``seed`` at
+        level >= 1."""
+        sp = getattr(req, "sampling_params", None)
+        view: dict = {"request_id": self._request_id(req)}
+        if self.log_requests_level >= 1:
+            view["seed"] = getattr(sp, "seed", None) if sp is not None else None
+        return view
+
+    def _batch_record(self, reqs: list) -> tuple:
+        """Build the ``rid`` / ``obj`` for one forward call"""
+        rids = [self._request_id(req) or "unknown" for req in reqs]
+        if len(reqs) == 1:
+            # Single request: scalar rid + flat dict obj (id + config).
+            req = reqs[0]
+            obj = {"request_id": self._request_id(req), **self._config_view(req)}
+            return rids[0], obj
+
+        shared_views = [self._config_view(req, drop_seed=True) for req in reqs]
+        if all(view == shared_views[0] for view in shared_views):
+            obj = {
+                **shared_views[0],
+                "outputs": [self._per_request_view(req) for req in reqs],
+            }
+        else:
+            # Configs genuinely differ: list each request's full payload verbatim.
+            obj = [
+                {"request_id": self._request_id(req), **self._config_view(req)}
+                for req in reqs
+            ]
+        return rids, obj
+
+    def _loggable(self, req: Any) -> bool:
+        """Whether ``req`` should be recorded: logging is on, it's a real
+        generation request (control messages -- LoRA / weight / stats / shutdown
+        -- have no ``sampling_params`` and are skipped), and it's not a warmup."""
+        return (
+            self.log_requests
+            and getattr(req, "sampling_params", None) is not None
+            and not getattr(req, "is_warmup", False)
+        )
+
+    def _logged_reqs(self, batch: Any) -> list:
+        """Normalize ``batch`` to a list and drop control / warmup requests."""
+        reqs = batch if isinstance(batch, (list, tuple)) else [batch]
+        return [r for r in reqs if self._loggable(r)]
+
+    def log_received_request(self, batch: Any) -> None:
+        reqs = self._logged_reqs(batch)
+        if not reqs:
+            return
+
+        rid, obj = self._batch_record(reqs)
+        max_length = self._max_length
+
+        if self.log_requests_format == "json":
+            log_json(
+                self.targets,
+                "request.received",
+                {"rid": rid, "obj": _transform_data_for_logging(obj, max_length)},
+            )
+        else:
+            self._emit(
+                f"Receive: obj={_dataclass_to_string_truncated(obj, max_length)}"
+            )
+
+    def log_finished_request(self, batch: Any, result: Any) -> None:
+        reqs = self._logged_reqs(batch)
+        if not reqs:
+            return
+
+        out = self._result_view(result)
+        e2e_latency_ms = out["meta_info"]["e2e_latency"] * 1000
+        if self.log_exceeded_ms > 0 and e2e_latency_ms < self.log_exceeded_ms:
+            return
+
+        rid, obj = self._batch_record(reqs)
+        max_length = self._max_length
+
+        if self.log_requests_format == "json":
+            log_json(
+                self.targets,
+                "request.finished",
+                {
+                    "rid": rid,
+                    "obj": _transform_data_for_logging(obj, max_length),
+                    "out": _transform_data_for_logging(out, max_length),
+                },
+            )
+        else:
+            self._emit(
+                f"Finish: obj={_dataclass_to_string_truncated(obj, max_length)}"
+                f", out={_dataclass_to_string_truncated(out, max_length)}"
+            )

python/sglang/multimodal_gen/runtime/utils/torch_compile.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.srt.utils.common import get_compiler_backend
+
+
+def maybe_enable_inductor_compute_comm_overlap() -> None:
+    try:
+        import torch._inductor.config as _inductor_cfg
+
+        _inductor_cfg.reorder_for_compute_comm_overlap = True
+    except ImportError:
+        pass
+
+
+def build_torch_compile_kwargs(*, mode: str | None) -> dict[str, object]:
+    compile_kwargs: dict[str, object] = {"fullgraph": False, "dynamic": None}
+    if current_platform.is_npu():
+        compile_kwargs["backend"] = get_compiler_backend()
+        compile_kwargs["dynamic"] = False
+    elif mode is not None:
+        compile_kwargs["mode"] = mode
+    return compile_kwargs
+
+
+def resolve_torch_compile_mode(
+    *env_names: str,
+    config: object | None = None,
+    default: str,
+) -> str:
+    for env_name in env_names:
+        mode = os.environ.get(env_name)
+        if mode:
+            return mode
+    mode = getattr(config, "torch_compile_mode", None)
+    if mode:
+        return mode
+    return default
+
+
+@dataclass
+class CompiledModuleRegistry:
+    module_ids: set[int] = field(default_factory=set)
+
+    def is_compiled(self, module: nn.Module) -> bool:
+        return id(module) in self.module_ids
+
+    def compile_once(
+        self,
+        module: nn.Module,
+        *,
+        compile_kwargs: dict[str, object],
+    ) -> bool:
+        module_id = id(module)
+        if module_id in self.module_ids:
+            return False
+        module.compile(**compile_kwargs)
+        self.module_ids.add(module_id)
+        return True
+
+
+class CallableModule(nn.Module):
+    """Module wrapper for compiling non-forward callables with module.compile"""
+
+    def __init__(self, fn: Callable[..., Any]) -> None:
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, *args, **kwargs):
+        return self.fn(*args, **kwargs)
+
+
+@dataclass
+class ActiveTargetCompiledCallable:
+    """Cache one compiled callable module for the currently active target object"""
+
+    target_id: int | None = None
+    compiled_module: CallableModule | None = None
+
+    def get_or_compile(
+        self,
+        target: object,
+        fn: Callable[..., Any],
+        *,
+        compile_kwargs: dict[str, object],
+    ) -> Callable[..., Any]:
+        target_id = id(target)
+        if self.target_id == target_id and self.compiled_module is not None:
+            return self.compiled_module
+
+        module = CallableModule(fn)
+        module.compile(**compile_kwargs)
+        self.target_id = target_id
+        self.compiled_module = module
+        return module

python/sglang/multimodal_gen/runtime/utils/trace_wrapper.py

@@ -0,0 +1,77 @@
+"""Context-manager wrappers around sglang.srt.observability.trace for diffusion tracing.
+
+All tracing helpers for the multimodal_gen subsystem are consolidated here so
+that call sites can use simple ``with`` statements instead of manual
+start/end bookkeeping.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+
+DIFFUSION_TRACE_MODULE = "diffusion"
+
+
+@dataclass(frozen=True)
+class DiffStageConfig:
+    """A named trace stage with a default nesting level."""
+
+    stage_name: str
+    level: int = 0
+
+
+class DiffStage:
+    """Named trace stages for the diffusion pipeline."""
+
+    SCHEDULER_DISPATCH = DiffStageConfig("scheduler_dispatch", level=1)
+    GPU_FORWARD = DiffStageConfig("gpu_forward", level=2)
+
+
+def init_diffusion_tracing(server_args, thread_label: str):
+    if not server_args.enable_trace:
+        return
+
+    from sglang.srt.observability.trace import (
+        process_tracing_init,
+        trace_set_thread_info,
+    )
+
+    # srt owns TraceReqContext and filters spans through its trace_modules list
+    process_tracing_init(
+        server_args.otlp_traces_endpoint,
+        "sglang-diffusion",
+        trace_modules=DIFFUSION_TRACE_MODULE,
+    )
+    trace_set_thread_info(thread_label)
+
+
+@contextmanager
+def trace_req(trace_ctx):
+    """Ensure ``trace_req_finish()`` is called when a request scope exits.
+
+    Usage::
+
+        with trace_req(batch.trace_ctx):
+            ...
+    """
+    try:
+        yield trace_ctx
+    finally:
+        trace_ctx.trace_req_finish()
+
+
+@contextmanager
+def trace_slice(trace_ctx, stage: DiffStageConfig, **kwargs):
+    """Context manager for a single trace slice (span).
+
+    Usage::
+
+        with trace_slice(req.trace_ctx, DiffStage.GPU_FORWARD):
+            result = pipeline.forward(req, server_args)
+    """
+    trace_ctx.trace_slice_start(stage.stage_name, level=stage.level)
+    try:
+        yield trace_ctx
+    finally:
+        trace_ctx.trace_slice_end(stage.stage_name, level=stage.level, **kwargs)

python/sglang/multimodal_gen/runtime/utils/vision.py

@@ -0,0 +1,297 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import tempfile
+from collections.abc import Callable
+from io import BytesIO
+from urllib.parse import unquote, urlparse
+
+import imageio
+import numpy as np
+import PIL.Image
+import PIL.ImageOps
+import requests
+import torch
+from packaging import version
+
+from sglang.srt.utils.common import get_image_bytes as srt_get_image_bytes
+
+if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"):
+    PIL_INTERPOLATION = {
+        "linear": PIL.Image.Resampling.BILINEAR,
+        "bilinear": PIL.Image.Resampling.BILINEAR,
+        "bicubic": PIL.Image.Resampling.BICUBIC,
+        "lanczos": PIL.Image.Resampling.LANCZOS,
+        "nearest": PIL.Image.Resampling.NEAREST,
+    }
+else:
+    PIL_INTERPOLATION = {
+        "linear": PIL.Image.LINEAR,
+        "bilinear": PIL.Image.BILINEAR,
+        "bicubic": PIL.Image.BICUBIC,
+        "lanczos": PIL.Image.LANCZOS,
+        "nearest": PIL.Image.NEAREST,
+    }
+
+
+def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
+    r"""
+    Convert a PIL image or a list of PIL images to NumPy arrays.
+
+    Args:
+        images (`PIL.Image.Image` or `List[PIL.Image.Image]`):
+            The PIL image or list of images to convert to NumPy format.
+
+    Returns:
+        `np.ndarray`:
+            A NumPy array representation of the images.
+    """
+    if not isinstance(images, list):
+        images = [images]
+    images = [np.array(image).astype(np.float32) / 255.0 for image in images]
+    images_arr: np.ndarray = np.stack(images, axis=0)
+
+    return images_arr
+
+
+def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
+    r"""
+    Convert a NumPy image to a PyTorch tensor.
+
+    Args:
+        images (`np.ndarray`):
+            The NumPy image array to convert to PyTorch format.
+
+    Returns:
+        `torch.Tensor`:
+            A PyTorch tensor representation of the images.
+    """
+    if images.ndim == 3:
+        images = images[..., None]
+
+    images = torch.from_numpy(images.transpose(0, 3, 1, 2))
+    return images
+
+
+def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
+    r"""
+    Normalize an image array to [-1,1].
+
+    Args:
+        images (`np.ndarray` or `torch.Tensor`):
+            The image array to normalize.
+
+    Returns:
+        `np.ndarray` or `torch.Tensor`:
+            The normalized image array.
+    """
+    return 2.0 * images - 1.0
+
+
+# adapted from diffusers.utils import load_image
+def load_image(
+    image: str | bytes | PIL.Image.Image,
+    convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None,
+) -> PIL.Image.Image:
+    """
+    Loads `image` to a PIL Image.
+
+    Args:
+        image (`str` or `PIL.Image.Image`):
+            The image to convert to the PIL Image format.
+        convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*):
+            A conversion method to apply to the image after loading it. When set to `None` the image will be converted
+            "RGB".
+    """
+    if isinstance(image, (str, bytes)):
+        if isinstance(image, str) and os.path.isfile(image):
+            image = PIL.Image.open(image)
+        else:
+            # in-memory loading path
+            image = PIL.Image.open(BytesIO(srt_get_image_bytes(image)))
+    elif isinstance(image, PIL.Image.Image):
+        image = image
+    else:
+        raise ValueError(
+            "Incorrect format used for the image. Should be bytes, a URL, a local path, base64/data URL, or a PIL image."
+        )
+
+    image = PIL.ImageOps.exif_transpose(image)
+
+    if convert_method is not None:
+        image = convert_method(image)
+    else:
+        image = image.convert("RGB")
+
+    return image
+
+
+# adapted from diffusers.utils import load_video
+def load_video(
+    video: str,
+    convert_method: (
+        Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]] | None
+    ) = None,
+) -> list[PIL.Image.Image]:
+    """
+    Loads `video` to a list of PIL Image.
+    Args:
+        video (`str`):
+            A URL or Path to a video to convert to a list of PIL Image format.
+        convert_method (Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]], *optional*):
+            A conversion method to apply to the video after loading it. When set to `None` the images will be converted
+            to "RGB".
+    Returns:
+        `List[PIL.Image.Image]`:
+            The video as a list of PIL images.
+    """
+    is_url = video.startswith("http://") or video.startswith("https://")
+    is_file = os.path.isfile(video)
+    was_tempfile_created = False
+
+    if not (is_url or is_file):
+        raise ValueError(
+            f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path."
+        )
+
+    if is_url:
+        response = requests.get(video, stream=True)
+        if response.status_code != 200:
+            raise ValueError(
+                f"Failed to download video. Status code: {response.status_code}"
+            )
+
+        parsed_url = urlparse(video)
+        file_name = os.path.basename(unquote(parsed_url.path))
+
+        suffix = os.path.splitext(file_name)[1] or ".mp4"
+        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as temp_file:
+            video_path = temp_file.name
+            video_data = response.iter_content(chunk_size=8192)
+            for chunk in video_data:
+                temp_file.write(chunk)
+
+        video = video_path
+
+    pil_images = []
+    if video.endswith(".gif"):
+        gif = PIL.Image.open(video)
+        try:
+            while True:
+                pil_images.append(gif.copy())
+                gif.seek(gif.tell() + 1)
+        except EOFError:
+            pass
+
+    else:
+        try:
+            imageio.plugins.ffmpeg.get_exe()
+        except AttributeError:
+            raise AttributeError(
+                "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg"
+            ) from None
+
+        with imageio.get_reader(video) as reader:
+            # Read all frames
+            for frame in reader:
+                pil_images.append(PIL.Image.fromarray(frame))
+
+    if was_tempfile_created:
+        os.remove(video_path)
+
+    if convert_method is not None:
+        pil_images = convert_method(pil_images)
+
+    return pil_images
+
+
+def get_default_height_width(
+    image: PIL.Image.Image | np.ndarray | torch.Tensor,
+    vae_scale_factor: int,
+    height: int | None = None,
+    width: int | None = None,
+) -> tuple[int, int]:
+    r"""
+    Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`.
+
+    Args:
+        image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`):
+            The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it
+            should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch
+            tensor, it should have shape `[batch, channels, height, width]`.
+        height (`Optional[int]`, *optional*, defaults to `None`):
+            The height of the preprocessed image. If `None`, the height of the `image` input will be used.
+        width (`Optional[int]`, *optional*, defaults to `None`):
+            The width of the preprocessed image. If `None`, the width of the `image` input will be used.
+
+    Returns:
+        `Tuple[int, int]`:
+            A tuple containing the height and width, both resized to the nearest integer multiple of
+            `vae_scale_factor`.
+    """
+
+    if height is None:
+        if isinstance(image, PIL.Image.Image):
+            height = image.height
+        elif isinstance(image, torch.Tensor):
+            height = image.shape[2]
+        else:
+            height = image.shape[1]
+
+    if width is None:
+        if isinstance(image, PIL.Image.Image):
+            width = image.width
+        elif isinstance(image, torch.Tensor):
+            width = image.shape[3]
+        else:
+            width = image.shape[2]
+
+    width, height = (
+        x - x % vae_scale_factor for x in (width, height)
+    )  # resize to integer multiple of vae_scale_factor
+
+    return height, width
+
+
+def resize(
+    image: PIL.Image.Image | np.ndarray | torch.Tensor,
+    height: int,
+    width: int,
+    resize_mode: str = "default",  # "default", "fill", "crop"
+    resample: str = "lanczos",
+) -> PIL.Image.Image | np.ndarray | torch.Tensor:
+    """
+    Resize image.
+
+    Args:
+        image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`):
+            The image input, can be a PIL image, numpy array or pytorch tensor.
+        height (`int`):
+            The height to resize to.
+        width (`int`):
+            The width to resize to.
+        resize_mode (`str`, *optional*, defaults to `default`):
+            The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit
+            within the specified width and height, and it may not maintaining the original aspect ratio. If `fill`,
+            will resize the image to fit within the specified width and height, maintaining the aspect ratio, and
+            then center the image within the dimensions, filling empty with data from image. If `crop`, will resize
+            the image to fit within the specified width and height, maintaining the aspect ratio, and then center
+            the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
+            supported for PIL image input.
+
+    Returns:
+        `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
+            The resized image.
+    """
+    if resize_mode != "default" and not isinstance(image, PIL.Image.Image):
+        raise ValueError(
+            f"Only PIL image input is supported for resize_mode {resize_mode}"
+        )
+    assert isinstance(image, PIL.Image.Image)
+    if resize_mode == "default":
+        image = image.resize((width, height), resample=PIL_INTERPOLATION[resample])
+    else:
+        raise ValueError(f"resize_mode {resize_mode} is not supported")
+    return image

python/sglang/multimodal_gen/runtime/utils/weight_attrs.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/utils.py
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+
+def set_weight_attrs(
+    weight: torch.Tensor,
+    weight_attrs: dict[str, Any] | None,
+):
+    """Set attributes on a weight tensor without overwriting existing ones."""
+    if weight_attrs is None:
+        return
+    for key, value in weight_attrs.items():
+        assert not hasattr(weight, key), f"Overwriting existing tensor attribute: {key}"
+
+        if current_platform.is_tpu() and key == "weight_loader":
+            value = make_synced_weight_loader(value)
+        setattr(weight, key, value)
+
+
+def make_synced_weight_loader(original_weight_loader) -> Any:
+
+    def _synced_weight_loader(param, *args, **kwargs):
+        original_weight_loader(param, *args, **kwargs)
+        torch._sync(param)
+
+    return _synced_weight_loader