sgl-project--sglang/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0
"""
DiffGenerator module for sglang-diffusion.

This module provides a consolidated interface for generating images/videos using
diffusion models.
"""

import dataclasses
import multiprocessing as mp
import os
import time
from contextlib import ExitStack
from typing import Any, List, Union

from sglang.multimodal_gen.configs.sample.sampling_params import (
    DataType,
    SamplingParams,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import (
    GenerationResult,
    ListLorasReq,
    MergeLoraWeightsReq,
    SetLoraReq,
    ShutdownReq,
    UnmergeLoraWeightsReq,
    expand_request_outputs,
    format_lora_message,
    prepare_request,
    save_outputs,
)
from sglang.multimodal_gen.runtime.launch_server import launch_server
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.scheduler_client import sync_scheduler_client
from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
from sglang.multimodal_gen.runtime.server_warmup import (
    run_sync_client_warmup,
    should_run_explicit_client_warmup,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import (
    GREEN,
    RESET,
    init_logger,
    log_batch_completion,
    log_generation_timer,
)
from sglang.multimodal_gen.runtime.utils.trace_wrapper import (
    init_diffusion_tracing,
    trace_req,
)

logger = init_logger(__name__)

try:
    # Set the start method to 'spawn' to avoid CUDA errors in forked processes.
    # This must be done at the top level of the module, before any CUDA context
    # or other processes are initialized.
    mp.set_start_method("spawn", force=True)
except RuntimeError:
    # The start method can only be set once per program execution.
    pass


class DiffGenerator:
    """
    A unified class for generating images/videos using diffusion models.

    This class provides a simple interface for image/video generation with rich
    customization options, similar to popular frameworks like HF Diffusers.
    """

    def __init__(
        self,
        server_args: ServerArgs,
    ):
        """
        Initialize the generator.

        Args:
            server_args: The inference arguments
        """
        self.server_args = server_args
        self.port_args = PortArgs.from_server_args(server_args)

        # The executor is now a client to the Scheduler service
        self.local_scheduler_process: list[mp.Process] | None = None
        self.owns_scheduler_client: bool = False

    @classmethod
    def from_pretrained(
        cls,
        local_mode: bool = True,
        **kwargs,
    ) -> "DiffGenerator":
        """
        Create a DiffGenerator from a pretrained model.

        Priority level: Default pipeline config < User's pipeline config < User's kwargs
        """
        # If users also provide some kwargs, it will override the ServerArgs and PipelineConfig.

        if (server_args := kwargs.get("server_args", None)) is not None:
            if isinstance(server_args, ServerArgs):
                pass
            elif isinstance(server_args, dict):
                server_args = ServerArgs.from_kwargs(**server_args)
        else:
            server_args = ServerArgs.from_kwargs(**kwargs)

        return cls.from_server_args(server_args, local_mode=local_mode)

    @classmethod
    def from_server_args(
        cls, server_args: ServerArgs, local_mode: bool = True
    ) -> "DiffGenerator":
        """
        Create a DiffGenerator with the specified arguments.

        Args:
            server_args: The inference arguments

        Returns:
            The created DiffGenerator
        """
        instance = cls(
            server_args=server_args,
        )
        init_diffusion_tracing(server_args, "DiffGenerator")

        logger.info(f"Local mode: {local_mode}")
        if local_mode:
            instance.local_scheduler_process = instance._start_local_server_if_needed()
            instance.owns_scheduler_client = True
            instance._run_client_warmup_if_needed()
        else:
            # In remote mode, we just need to connect and check.
            sync_scheduler_client.initialize(server_args)
            instance._check_remote_scheduler()
            instance.owns_scheduler_client = True
        return instance

    def _start_local_server_if_needed(
        self,
    ) -> list[mp.Process]:
        """Check if a local server is running; if not, start it and return the process handles."""
        # First, we need a client to test the server. Initialize it temporarily.
        sync_scheduler_client.initialize(self.server_args)

        processes = launch_server(self.server_args, launch_http_server=False)

        return processes

    def _run_client_warmup_if_needed(self) -> None:
        if not should_run_explicit_client_warmup(self.server_args):
            return

        run_sync_client_warmup(self.server_args, sync_scheduler_client.forward)

    def _check_remote_scheduler(self):
        """Check if the remote scheduler is accessible."""
        if not sync_scheduler_client.ping():
            raise ConnectionError(
                f"Could not connect to remote scheduler at "
                f"{self.server_args.scheduler_endpoint} with `local mode` as False. "
                "Please ensure the server is running."
            )
        logger.info(
            f"Successfully connected to remote scheduler at "
            f"{self.server_args.scheduler_endpoint}."
        )

    @staticmethod
    def _resolve_image_paths_per_prompt(
        prompts: list[str], image_paths: str | list[str] | None
    ) -> list[str | list[str] | None]:
        if len(prompts) <= 1:
            return [image_paths]

        if not isinstance(image_paths, list) or len(image_paths) <= 1:
            return [image_paths for _ in prompts]

        if len(image_paths) != len(prompts):
            raise ValueError(
                "When using multiple prompts with multiple input images, "
                "provide either one shared image or exactly one image per prompt."
            )

        return [[image_path] for image_path in image_paths]

    def generate(
        self,
        sampling_params_kwargs: dict | None = None,
        external_trace_header: dict[str, str] | None = None,
    ) -> GenerationResult | list[GenerationResult] | None:
        """Generate image(s)/video(s) based on the given prompt(s).

        Returns a single GenerationResult for a single prompt, a list for
        multiple prompts, or None when every request failed.
        """
        # 1. prepare requests
        prompts = self._resolve_prompts(
            sampling_params_kwargs.get("prompt"),
            sampling_params_kwargs.get("prompt_path"),
        )
        user_output_file_name = sampling_params_kwargs.get("output_file_name")

        if len(prompts) > 1 and user_output_file_name is not None:
            raise ValueError(
                "Cannot use multiple prompts with a fixed output_file_name. "
                "Either remove --output-file-name or use a single prompt."
            )

        sampling_params_orig = SamplingParams.from_user_sampling_params_args(
            self.server_args.model_path,
            server_args=self.server_args,
            **sampling_params_kwargs,
        )

        request_groups: list[list[Req]] = []
        image_paths_per_prompt = self._resolve_image_paths_per_prompt(
            prompts, sampling_params_orig.image_path
        )

        for i, p in enumerate(prompts):
            sampling_params = dataclasses.replace(
                sampling_params_orig,
                prompt=p,
                output_file_name=user_output_file_name,
                image_path=image_paths_per_prompt[i],
            )
            # `dataclasses.replace` drops non-field attrs; restore
            # `_explicit_fields` so InputValidationStage honors user-supplied
            # width/height, and mark the keys overridden above as explicit.
            sampling_params._explicit_fields = getattr(
                sampling_params_orig, "_explicit_fields", set()
            ) | {"prompt", "output_file_name", "image_path"}
            sampling_params._set_output_file_name()
            req = prepare_request(
                server_args=self.server_args,
                sampling_params=sampling_params,
                external_trace_header=external_trace_header,
            )
            request_groups.append(
                expand_request_outputs(
                    req,
                    num_prompts=len(prompts),
                    prompt_index=i,
                )
            )

        results: list[GenerationResult] = []
        total_start_time = time.perf_counter()
        global_output_index = 0

        for requests in request_groups:
            try:
                timer_prompt = [req.prompt for req in requests]
                logger.info("Processing %d grouped request(s)", len(requests))
                with ExitStack() as stack:
                    for req in requests:
                        stack.enter_context(trace_req(req.trace_ctx))
                    timer = stack.enter_context(
                        log_generation_timer(logger, timer_prompt)
                    )
                    output_batch = self._send_to_scheduler_and_wait_for_response(
                        requests
                    )
                    if output_batch.error:
                        raise Exception(f"{output_batch.error}")

                    if (
                        output_batch.output is None
                        and output_batch.output_file_paths is None
                    ):
                        logger.error("Received empty output from scheduler")
                        continue

                    if requests[0].save_output and requests[0].return_file_paths_only:
                        output_file_paths = output_batch.output_file_paths or []
                        self._validate_output_count(
                            len(output_file_paths), len(requests)
                        )
                        for idx, path in enumerate(output_file_paths):
                            req = requests[idx]
                            results.append(
                                GenerationResult(
                                    **self._result_common(
                                        req, output_batch, timer.duration, idx
                                    ),
                                    prompt_index=global_output_index + idx,
                                    output_file_path=path,
                                )
                            )
                    elif requests[0].data_type == DataType.MESH:
                        output_file_paths = output_batch.output_file_paths or []
                        self._validate_output_count(
                            len(output_file_paths), len(requests)
                        )
                        for idx, sample in enumerate(output_file_paths):
                            req = requests[idx]
                            results.append(
                                GenerationResult(
                                    **self._result_common(
                                        req, output_batch, timer.duration, idx
                                    ),
                                    prompt_index=global_output_index + idx,
                                    output_file_path=sample,
                                )
                            )
                    else:
                        self._validate_output_count(
                            len(output_batch.output), len(requests)
                        )
                        samples_out: list[Any] = []
                        audios_out: list[Any] = []
                        frames_out: list[Any] = []
                        save_outputs(
                            output_batch.output,
                            requests[0].data_type,
                            requests[0].fps,
                            requests[0].save_output,
                            lambda idx: requests[idx].output_file_path(1, 0),
                            audio=output_batch.audio,
                            audio_sample_rate=output_batch.audio_sample_rate,
                            samples_out=samples_out,
                            audios_out=audios_out,
                            frames_out=frames_out,
                            output_compression=requests[0].output_compression,
                            enable_frame_interpolation=requests[
                                0
                            ].enable_frame_interpolation,
                            frame_interpolation_exp=requests[0].frame_interpolation_exp,
                            frame_interpolation_scale=requests[
                                0
                            ].frame_interpolation_scale,
                            frame_interpolation_model_path=requests[
                                0
                            ].frame_interpolation_model_path,
                            enable_upscaling=requests[0].enable_upscaling,
                            upscaling_model_path=requests[0].upscaling_model_path,
                            upscaling_scale=requests[0].upscaling_scale,
                        )

                        for idx in range(len(samples_out)):
                            req = requests[idx]
                            results.append(
                                GenerationResult(
                                    **self._result_common(
                                        req, output_batch, timer.duration, idx
                                    ),
                                    samples=samples_out[idx],
                                    frames=frames_out[idx],
                                    audio=audios_out[idx],
                                    prompt_index=global_output_index + idx,
                                    output_file_path=req.output_file_path(1, 0),
                                )
                            )
            except Exception as e:
                logger.error("Generation failed: %s", e, exc_info=True)
            finally:
                global_output_index += len(requests)

        total_gen_time = time.perf_counter() - total_start_time
        if self.server_args.batching_max_size > 1:
            log_batch_completion(
                logger,
                len(results),
                total_gen_time,
            )
        self._log_summary(results)

        if not results:
            return None
        return results[0] if len(results) == 1 else results

    def generate_action(
        self,
        sampling_params_kwargs: dict | None = None,
        external_trace_header: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        sampling_params_kwargs = sampling_params_kwargs or {}
        sampling_params = SamplingParams.from_user_sampling_params_args(
            self.server_args.model_path,
            server_args=self.server_args,
            **sampling_params_kwargs,
        )
        if sampling_params.data_type != DataType.ACTION:
            raise ValueError(
                f"generate_action requires an ACTION pipeline, got {sampling_params.data_type}"
            )

        req = prepare_request(
            server_args=self.server_args,
            sampling_params=sampling_params,
            external_trace_header=external_trace_header,
        )
        output_batch = self._send_to_scheduler_and_wait_for_response(req)
        if output_batch.error:
            raise RuntimeError(output_batch.error)
        if output_batch.output is None:
            raise RuntimeError("action policy returned no output")
        return output_batch.output[0]

    def _resolve_prompts(
        self,
        prompt: str | list[str] | None,
        prompt_path: str | None = None,
    ) -> list[str]:
        """Collect prompts from the argument or from a prompt file."""
        path = prompt_path or self.server_args.prompt_file_path
        if path is not None:
            if not os.path.exists(path):
                raise FileNotFoundError(f"Prompt text file not found: {path}")
            with open(path, encoding="utf-8") as f:
                prompts = [line.strip() for line in f if line.strip()]
            if not prompts:
                raise ValueError(f"No prompts found in file: {path}")
            logger.info("Found %d prompts in %s", len(prompts), path)
            return prompts

        if prompt is None:
            return [" "]
        if isinstance(prompt, str):
            return [prompt]
        return list(prompt)

    def _log_summary(self, results: list[GenerationResult]) -> None:
        if not results:
            return
        if self.server_args.warmup:
            total_duration_ms = results[0].metrics.get("total_duration_ms", 0)
            logger.info(
                f"Warmed-up request processed in {GREEN}%.2f{RESET} seconds (with warmup excluded)",
                total_duration_ms / 1000.0,
            )

        peak_memories = [r.peak_memory_mb for r in results if r.peak_memory_mb]
        if peak_memories:
            logger.info(
                f"Memory usage - Max peak: {max(peak_memories):.2f} MB, "
                f"Avg peak: {sum(peak_memories) / len(peak_memories):.2f} MB"
            )

    @staticmethod
    def _result_common(
        req: Req,
        output_batch: OutputBatch,
        generation_time: float,
        output_index: int | None = None,
    ) -> dict[str, Any]:
        metrics = output_batch.metrics
        if (
            output_index is not None
            and output_batch.metrics_list is not None
            and output_index < len(output_batch.metrics_list)
        ):
            metrics = output_batch.metrics_list[output_index]
        if req.data_type == DataType.ACTION:
            size = ("action",)
        else:
            size = (req.height, req.width, req.num_frames)
        return dict(
            prompt=req.prompt,
            size=size,
            generation_time=generation_time,
            peak_memory_mb=output_batch.peak_memory_mb,
            metrics=metrics.to_dict() if metrics else {},
            action=output_batch.action_pred,
            trajectory_latents=output_batch.trajectory_latents,
            trajectory_timesteps=output_batch.trajectory_timesteps,
            rollout_trajectory_data=output_batch.rollout_trajectory_data,
            trajectory_decoded=output_batch.trajectory_decoded,
        )

    @staticmethod
    def _validate_output_count(output_count: int, request_count: int) -> None:
        if output_count != request_count:
            raise RuntimeError(
                f"Expected {request_count} outputs, got {output_count} from scheduler"
            )

    def _send_to_scheduler_and_wait_for_response(self, batch: list[Req]) -> OutputBatch:
        """
        Sends a request to the scheduler and waits for a response.
        """
        return sync_scheduler_client.forward(batch)

    # LoRA
    def _send_lora_request(self, req: Any, success_msg: str, failure_msg: str):
        response = sync_scheduler_client.forward(req)
        if response.error is None:
            logger.info(success_msg)
            return response
        else:
            error_msg = response.error
            raise RuntimeError(f"{failure_msg}: {error_msg}")

    def set_lora(
        self,
        lora_nickname: Union[str, List[str]],
        lora_path: Union[str, None, List[Union[str, None]]] = None,
        target: Union[str, List[str]] = "all",
        strength: Union[float, List[float]] = 1.0,
        merge_mode: str | None = None,
    ) -> None:
        """
        Set LoRA adapter(s) for the specified transformer(s).
        Supports both single LoRA (backward compatible) and multiple LoRA adapters.

        Args:
            lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings.
            lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None.
            target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings.
                Valid values:
                - "all": Apply to all transformers (default)
                - "transformer": Apply only to the primary transformer (high noise for Wan2.2)
                - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2)
                - "critic": Apply only to the critic model
            strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
            merge_mode: Optional LoRA merge mode: "auto", "merge", or "dynamic".
        """
        req = SetLoraReq(
            lora_nickname=lora_nickname,
            lora_path=lora_path,
            target=target,
            strength=strength,
            merge_mode=merge_mode,
        )
        nickname_str, target_str, strength_str = format_lora_message(
            lora_nickname, target, strength
        )

        self._send_lora_request(
            req,
            f"Successfully set LoRA adapter(s): {nickname_str} (target: {target_str}, strength: {strength_str})",
            "Failed to set LoRA adapter",
        )

    def unmerge_lora_weights(self, target: str = "all") -> None:
        """
        Unmerge LoRA weights from the base model.

        Args:
            target: Which transformer(s) to unmerge.
        """
        req = UnmergeLoraWeightsReq(target=target)
        self._send_lora_request(
            req,
            f"Successfully unmerged LoRA weights (target: {target})",
            "Failed to unmerge LoRA weights",
        )

    def merge_lora_weights(self, target: str = "all", strength: float = 1.0) -> None:
        """
        Merge LoRA weights into the base model.

        Args:
            target: Which transformer(s) to merge.
            strength: LoRA strength for merge, default 1.0.
        """
        req = MergeLoraWeightsReq(target=target, strength=strength)
        self._send_lora_request(
            req,
            f"Successfully merged LoRA weights (target: {target}, strength: {strength})",
            "Failed to merge LoRA weights",
        )

    def list_loras(self) -> dict:
        """List loaded LoRA adapters and current application status per module."""
        output = self._send_lora_request(
            req=ListLorasReq(),
            success_msg="Successfully listed LoRA adapters",
            failure_msg="Failed to list LoRA adapters",
        )
        # _send_lora_request already raises on error, so output.error is always None here
        return output.output or {}

    def _ensure_lora_state(
        self,
        lora_path: str | None,
        lora_nickname: str | None = None,
        merge_lora: bool = True,
    ) -> None:
        """
        Ensure the LoRA state matches the desired configuration.

        Note: This method does not cache client-side state. The server handles
        idempotent operations, so redundant calls are safe but may have minor overhead.
        """
        if lora_path is None:
            # Unmerge all LoRA weights when no lora_path is provided
            self.unmerge_lora_weights()
            return

        lora_nickname = lora_nickname or self.server_args.lora_nickname

        # Set the LoRA adapter (server handles idempotent logic)
        self.set_lora(lora_nickname, lora_path)

        # Merge or unmerge based on the merge_lora flag
        if merge_lora:
            self.merge_lora_weights()
        else:
            self.unmerge_lora_weights()

    def generate_with_lora(
        self,
        prompt: str | list[str] | None = None,
        sampling_params: SamplingParams | None = None,
        *,
        lora_path: str | None = None,
        lora_nickname: str | None = None,
        merge_lora: bool = True,
        **kwargs,
    ):
        self._ensure_lora_state(
            lora_path=lora_path, lora_nickname=lora_nickname, merge_lora=merge_lora
        )
        return self.generate(
            sampling_params_kwargs=dict(
                prompt=prompt,
                sampling_params=sampling_params,
                **kwargs,
            )
        )

    def shutdown(self):
        """
        Shutdown the generator.
        If in local mode, it also shuts down the scheduler server.
        """
        # sends the shutdown command to the server
        if self.local_scheduler_process and self.owns_scheduler_client:
            try:
                sync_scheduler_client.forward(ShutdownReq(), timeout_ms=5000)
            except Exception:
                pass

        if self.local_scheduler_process:
            for process in self.local_scheduler_process:
                process.join(timeout=10)
                if process.is_alive():
                    logger.warning(
                        f"Local worker {process.name} did not terminate gracefully, forcing."
                    )
                    process.terminate()
                    process.join(timeout=1)
                    if process.is_alive():
                        process.kill()
                        process.join(timeout=1)
            self.local_scheduler_process = None

        if self.owns_scheduler_client:
            sync_scheduler_client.close()
            self.owns_scheduler_client = False

    def _force_shutdown_local_processes(self) -> None:
        local_scheduler_process = getattr(self, "local_scheduler_process", None)
        log = globals().get("logger")
        if local_scheduler_process:
            for process in local_scheduler_process:
                if process.is_alive():
                    if log is not None:
                        log.warning(
                            f"Local worker {process.name} did not terminate gracefully, forcing."
                        )
                    process.terminate()
            for process in local_scheduler_process:
                process.join(timeout=1)
                if process.is_alive():
                    if log is not None:
                        log.warning(
                            f"Local worker {process.name} did not terminate after terminate(), killing."
                        )
                    process.kill()
                    process.join(timeout=1)
            self.local_scheduler_process = None

        if getattr(self, "owns_scheduler_client", False):
            try:
                client = globals().get("sync_scheduler_client")
                if client is not None:
                    client.close()
            finally:
                self.owns_scheduler_client = False

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.shutdown()

    def __del__(self):
        owns_scheduler_client = bool(getattr(self, "owns_scheduler_client", False))
        local_scheduler_process = getattr(self, "local_scheduler_process", None)
        log = globals().get("logger")
        if owns_scheduler_client:
            if log is not None:
                log.warning(
                    "Generator was garbage collected without being shut down. "
                    "Forcing local server and client cleanup."
                )
            self._force_shutdown_local_processes()
        elif local_scheduler_process:
            if log is not None:
                log.warning(
                    "Generator was garbage collected without being shut down. "
                    "Forcing local server cleanup."
                )
            self._force_shutdown_local_processes()