[diffusion] chore: minor code cleanups (#15190)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/multimodal_gen/apps/webui/main.py

@@ -5,8 +5,10 @@ from sglang.multimodal_gen.configs.sample.sampling_params import (
     DataType,
     SamplingParams,
 )
-from sglang.multimodal_gen.runtime.entrypoints.openai.utils import post_process_sample
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
-from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+    post_process_sample,
+    prepare_request,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client
 

python/sglang/multimodal_gen/docs/profiling.md

@@ -45,7 +45,7 @@ sglang generate \
 By default, trace files are saved in the ./logs/ directory. The exact output file path will be shown in the console output, for example:
 
 ```bash
-[mm-dd hh:mm:ss] Saving profiler traces to: /sgl-workspace/sglang/logs/mocked_fake_id_for_offline_generate-5_steps-global-rank0.trace.json.gz
+[mm-dd hh:mm:ss] Saved profiler traces to: /sgl-workspace/sglang/logs/mocked_fake_id_for_offline_generate-5_steps-global-rank0.trace.json.gz
 ```
 {request_id}-{num_steps}_steps-global-rank{rank}.trace.json.gz
 ```

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -13,22 +13,18 @@ import os
 import time
 from typing import Any
 
-import imageio
 import numpy as np
-import torch
-import torchvision
-from einops import rearrange
 
-from sglang.multimodal_gen.configs.sample.sampling_params import (
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
-    DataType,
-    SamplingParams,
-)
 from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
     MergeLoraWeightsReq,
     SetLoraReq,
     UnmergeLoraWeightsReq,
 )
-from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+from sglang.multimodal_gen.runtime.entrypoints.utils import (
+    post_process_sample,
+    prepare_request,
+)
 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
@@ -39,7 +35,6 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
     log_batch_completion,
     log_generation_timer,
     suppress_loggers,
-    suppress_other_loggers,
 )
 
 suppress_loggers(["imageio", "imageio_ffmpeg", "PIL", "PIL_Image"])
@@ -162,49 +157,6 @@ class DiffGenerator:
             f"{self.server_args.scheduler_endpoint()}."
         )
 
-    def post_process_sample(
-        self,
-        sample: torch.Tensor,
-        data_type: DataType,
-        fps: int,
-        save_output: bool = True,
-        save_file_path: str = None,
-    ):
-        """
-        Process a single sample output and save output if necessary
-        """
-        # Process outputs
-        if sample.dim() == 3:
-            # for images, dim t is missing
-            sample = sample.unsqueeze(1)
-        sample = rearrange(sample, "c t h w -> t c h w")
-        frames = []
-        # TODO: this can be batched
-        for x in sample:
-            x = torchvision.utils.make_grid(x, nrow=6)
-            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-            frames.append((x * 255).numpy().astype(np.uint8))
-
-        # Save outputs if requested
-        if save_output:
-            if save_file_path:
-                os.makedirs(os.path.dirname(save_file_path), exist_ok=True)
-                with suppress_other_loggers():
-                    if data_type == DataType.VIDEO:
-                        imageio.mimsave(
-                            save_file_path,
-                            frames,
-                            fps=fps,
-                            format=data_type.get_default_extension(),
-                        )
-                    else:
-                        imageio.imwrite(save_file_path, frames[0])
-                logger.info("Saved output to %s", save_file_path)
-            else:
-                logger.warning("No output path provided, output not saved")
-
-        return frames
-
     def generate(
         self,
         sampling_params_kwargs: dict | None = None,
@@ -280,7 +232,7 @@ class DiffGenerator:
                         continue
                     for output_idx, sample in enumerate(output_batch.output):
                         num_outputs = len(output_batch.output)
-                        frames = self.post_process_sample(
+                        frames = post_process_sample(
                             sample,
                             fps=req.fps,
                             save_output=req.save_output,

python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py

@@ -4,14 +4,9 @@ import os
 import time
 from typing import Optional
 
-import imageio
-import numpy as np
-import torch
-import torchvision
-from einops import rearrange
 from fastapi import UploadFile
 
-from sglang.multimodal_gen.configs.sample.sampling_params import DataType
+from sglang.multimodal_gen.runtime.entrypoints.utils import post_process_sample
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
     init_logger,
     log_batch_completion,
@@ -38,47 +33,6 @@ class UnmergeLoraWeightsReq:
     target: str = "all"  # "all", "transformer", "transformer_2", "critic"
 
 
-def post_process_sample(
-    sample: torch.Tensor,
-    data_type: DataType,
-    fps: int,
-    save_output: bool = True,
-    save_file_path: str = None,
-):
-    """
-    Process sample output and save video if necessary
-    """
-    # Process outputs
-    if sample.dim() == 3:
-        # for images, dim t is missing
-        sample = sample.unsqueeze(1)
-    videos = rearrange(sample, "c t h w -> t c h w")
-    frames = []
-    for x in videos:
-        x = torchvision.utils.make_grid(x, nrow=6)
-        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-        frames.append((x * 255).numpy().astype(np.uint8))
-
-    # Save outputs if requested
-    if save_output:
-        if save_file_path:
-            os.makedirs(os.path.dirname(save_file_path), exist_ok=True)
-            if data_type == DataType.VIDEO:
-                imageio.mimsave(
-                    save_file_path,
-                    frames,
-                    fps=fps,
-                    format=data_type.get_default_extension(),
-                )
-            else:
-                imageio.imwrite(save_file_path, frames[0])
-            logger.info(f"Saved output to {save_file_path}")
-        else:
-            logger.info(f"No output path provided, output not saved")
-
-    return frames
-
-
 def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]:
     try:
         parts = size.lower().replace(" ", "").split("x")
@@ -87,7 +41,6 @@ def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]:
         w, h = int(parts[0]), int(parts[1])
         return w, h
     except Exception:
-        # Fallback to default portrait 720x1280
         return None, None
 
 

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

@@ -9,11 +9,21 @@ diffusion models.
 """
 
 import dataclasses
+import os
 
-from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+import imageio
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger
 from sglang.multimodal_gen.utils import shallow_asdict
 
 logger = init_logger(__name__)
@@ -44,3 +54,45 @@ def prepare_request(
         )
 
     return req
+
+
+def post_process_sample(
+    sample: torch.Tensor,
+    data_type: DataType,
+    fps: int,
+    save_output: bool = True,
+    save_file_path: str = None,
+):
+    """
+    Process sample output and save video if necessary
+    """
+    # Process outputs
+    if sample.dim() == 3:
+        # for images, dim t is missing
+        sample = sample.unsqueeze(1)
+    videos = rearrange(sample, "c t h w -> t c h w")
+    frames = []
+    # TODO: this can be batched
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=6)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+        frames.append((x * 255).numpy().astype(np.uint8))
+
+    # Save outputs if requested
+    if save_output:
+        if save_file_path:
+            os.makedirs(os.path.dirname(save_file_path), exist_ok=True)
+            if data_type == DataType.VIDEO:
+                imageio.mimsave(
+                    save_file_path,
+                    frames,
+                    fps=fps,
+                    format=data_type.get_default_extension(),
+                )
+            else:
+                imageio.imwrite(save_file_path, frames[0])
+            logger.info(f"Saved output to {CYAN}{save_file_path}{RESET}")
+        else:
+            logger.info(f"No output path provided, output not saved")
+
+    return frames

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -128,7 +128,12 @@ class ComponentLoader(ABC):
                 component_model_path, server_args, module_name
             )
             source = "customized"
-        except Exception as _e:
+        except Exception as e:
+            if "Unsupported model architecture" in str(e):
+                logger.info(
+                    f"Module: {module_name} doesn't have a customized version yet, using native version"
+                )
+            else:
                 traceback.print_exc()
                 logger.error(
                     f"Error while loading customized {module_name}, falling back to native version"

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -33,9 +33,6 @@ from sglang.multimodal_gen.runtime.utils.perf_logger import (
 
 logger = init_logger(__name__)
 
-CYAN = "\033[1;36m"
-RESET = "\033[0;0m"
-
 
 class GPUWorker:
     """

python/sglang/multimodal_gen/runtime/models/registry.py

@@ -316,8 +316,9 @@ class _ModelRegistry:
         normalized_arch = []
         for arch in architectures:
             if arch not in self.registered_models:
+                registered_models = list(self.registered_models.keys())
                 raise Exception(
-                    f"Unsupported model architecture: {arch}. Registered architectures: {self.registered_models=}"
+                    f"Unsupported model architecture: {arch}. Registered architectures: {registered_models}"
                 )
             normalized_arch.append(arch)
         return normalized_arch

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -1114,10 +1114,8 @@ class DenoisingStage(PipelineStage):
             A tqdm progress bar.
         """
         local_rank = get_world_group().local_rank
-        if local_rank == 0:
+        disable = local_rank != 0
-            return tqdm(iterable=iterable, total=total)
+        return tqdm(iterable=iterable, total=total, disable=disable)
-        else:
-            return tqdm(iterable=iterable, total=total, disable=True)
 
     def rescale_noise_cfg(
         self, noise_cfg, noise_pred_text, guidance_rescale=0.0

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

@@ -24,6 +24,8 @@ SGLANG_DIFFUSION_LOGGING_CONFIG_PATH = envs.SGLANG_DIFFUSION_LOGGING_CONFIG_PATH
 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"
@@ -484,10 +486,6 @@ def log_generation_timer(
             total_requests,
             prompt[:100],
         )
-    else:
-        max_len = 100
-        suffix = "..." if len(prompt) > max_len else ""
-        logger.info(f"Processing prompt: {prompt[:100]}{suffix}")
 
     timer = GenerationTimer()
     timer.start_time = time.perf_counter()

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

@@ -2,7 +2,7 @@ import os
 
 import torch
 
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger
 
 logger = init_logger(__name__)
 
@@ -127,7 +127,7 @@ class SGLDiffusionProfiler:
                     f"{self.request_id}-{sanitized_profile_mode_id}-global-rank{dump_rank}.trace.json.gz",
                 )
             )
-            logger.info(f"Saving profiler traces to: {trace_path}")
             self.profiler.export_chrome_trace(trace_path)
+            logger.info(f"Saved profiler traces to: {CYAN}{trace_path}{RESET}")
         except Exception as e:
-            logger.error(f"Failed to export trace: {e}")
+            logger.error(f"Failed to save trace: {e}")