[diffusion] app: support webui (#14961)

Co-authored-by: Mick <mickjagger19@icloud.com>
2025-12-14 21:07:36 +08:00
parent e3f51e823e
commit 6f0c77d7f8
6 changed files with 257 additions and 11 deletions
--- a/python/sglang/multimodal_gen/apps/README.md
+++ b/python/sglang/multimodal_gen/apps/README.md
@@ -0,0 +1,53 @@
+# SGLang Diffusion WebUI User Guide
+
+SGLang Diffusion WebUI provides an intuitive Gradio-based interface for image and video generation, supporting parameter tuning and real-time previews.
+
+## Prerequisites
+The WebUI runs on Gradio. To get started, install Gradio first:
+```bash
+pip install gradio==6.1.0
+```
+
+## Launch WebUI Service
+SGLang Diffusion now includes an integrated WebUI. Simply add the `--webui` parameter when starting the service.
+
+### Launch Text-to-Image Service
+```bash
+SERVER_ARGS=(
+  --model-path black-forest-labs/FLUX.1-dev
+  --num-gpus 2
+)
+WEBUI_PORT=2333
+sglang serve "${SERVER_ARGS[@]}" --webui --webui-port ${WEBUI_PORT}
+```
+### Launch Text-to-Video Service
+```bash
+SERVER_ARGS=(
+  --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers
+  --num-gpus 2
+)
+WEBUI_PORT=2333
+sglang serve "${SERVER_ARGS[@]}" --webui --webui-port ${WEBUI_PORT}
+```
+
+## Port Forwarding
+Once the WebUI service is running, you need to use **SSH port forwarding** to securely access the remote service from your local machine.
+
+In most cases: Your IDE (like VS Code, Cursor, etc.) can handle this automatically. Check your IDE's remote development or port forwarding features. Otherwise, execute this command manually.
+
+```bash
+ssh -L ${WEBUI_PORT}:localhost:${WEBUI_PORT} user_name@machine_name
+```
+Learn more about port forwarding: [Port Forwarding](https://en.wikipedia.org/wiki/Port_forwarding).
+
+
+## Interface Instructions
+
+1. Task mode is automatically determined by the `num_frames` parameter:
+  - num_frames = 1: Text-to-Image mode
+  - num_frames > 1: Text-to-Video mode
+2. After generation, manually click:
+  - Image output: View generated images
+  - Video output: View generated videos
+
+Once launched, access the interface at `http://localhost:${WEBUI_PORT}` in your browser.
--- a/python/sglang/multimodal_gen/apps/webui/init.py
+++ b/python/sglang/multimodal_gen/apps/webui/init.py
@@ -0,0 +1,3 @@
+from .main import run_sgl_diffusion_webui
+
+__all__ = ["run_sgl_diffusion_webui"]
--- a/python/sglang/multimodal_gen/apps/webui/main.py
+++ b/python/sglang/multimodal_gen/apps/webui/main.py
@@ -0,0 +1,157 @@
+import argparse
+import os
+
+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 prepare_request
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client
+
+
+def add_webui_args(parser: argparse.ArgumentParser):
+    """Add the arguments for the generate command."""
+    parser = ServerArgs.add_cli_args(parser)
+    parser = SamplingParams.add_cli_args(parser)
+    return parser
+
+
+def run_sgl_diffusion_webui(server_args: ServerArgs):
+    # import gradio in function to avoid CI crash
+    import gradio as gr
+
+    # init client
+    sync_scheduler_client.initialize(server_args)
+
+    # server_args will be reused in gradio_generate function
+    def gradio_generate(
+        prompt,
+        negative_prompt,
+        seed,
+        num_frames,
+        frames_per_second,
+        width,
+        height,
+        num_inference_steps,
+        guidance_scale,
+        enable_teacache,
+    ):
+        """
+        NOTE: The input and output of function which is called by gradio button must be gradio components
+        So we use global variable sampling_params_kwargs to avoid pass this param, because gradio does not support this.
+        return [ np.ndarray, None ] | [None, np.ndarray]
+        """
+        sampling_params_kwargs = dict(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            seed=seed,
+            num_frames=num_frames,
+            fps=frames_per_second,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            enable_teacache=enable_teacache,
+        )
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            server_args.model_path,
+            server_args=server_args,
+            **sampling_params_kwargs,
+        )
+        batch = prepare_request(
+            server_args=server_args,
+            sampling_params=sampling_params,
+        )
+        result = sync_scheduler_client.forward([batch])
+        save_file_path = str(os.path.join(batch.output_path, batch.output_file_name))
+        frames = post_process_sample(
+            result.output[0],
+            batch.data_type,
+            batch.fps,
+            batch.save_output,
+            save_file_path,
+        )
+        if batch.data_type == DataType.VIDEO:
+            # gradio video need video path to show video
+            return None, save_file_path
+        else:
+            return frames[0], None
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# 🚀 SGLang Diffusion Application")
+        launched_model = gr.Textbox(label="Model", value=server_args.model_path)
+
+        with gr.Row():
+            with gr.Column(scale=4):
+                prompt = gr.Textbox(label="Prompt", value="A curious raccoon")
+                negative_prompt = gr.Textbox(
+                    label="Negative_prompt",
+                    value="Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
+                )
+            with gr.Column(scale=1):
+                seed = gr.Number(label="seed", precision=0, value=1234)
+                run_btn = gr.Button("Generate", variant="primary", size="lg")
+
+        with gr.Row():
+            with gr.Column():
+                num_frames = gr.Slider(
+                    minimum=1, maximum=161, value=81, step=1, label="num_frames"
+                )
+                frames_per_second = gr.Slider(
+                    minimum=4, maximum=60, value=16, step=1, label="frames_per_second"
+                )
+                width = gr.Number(label="width", precision=0, value=720)
+                height = gr.Number(label="height", precision=0, value=480)
+                num_inference_steps = gr.Slider(
+                    minimum=0, maximum=50, value=20, step=1, label="num_inference_steps"
+                )
+                guidance_scale = gr.Slider(
+                    minimum=0.0, maximum=10, value=5, step=0.01, label="guidance_scale"
+                )
+                enable_teacache = gr.Checkbox(label="enable_teacache", value=False)
+
+            with gr.Tabs() as tabs:
+                with gr.TabItem("Image output", id=1):
+                    image_out = gr.Image(
+                        label="Generated Image",
+                    )
+                with gr.TabItem("Video output", id=2):
+                    video_out = gr.Video(label="Generated Video")
+
+        run_btn.click(
+            fn=gradio_generate,
+            inputs=[
+                prompt,
+                negative_prompt,
+                seed,
+                num_frames,
+                frames_per_second,
+                width,
+                height,
+                num_inference_steps,
+                guidance_scale,
+                enable_teacache,
+            ],
+            outputs=[image_out, video_out],
+        )
+
+        _, local_url, _ = demo.launch(
+            server_port=server_args.webui_port,
+            quiet=True,
+            prevent_thread_lock=True,
+        )
+
+        # print banner
+        delimiter = "=" * 80
+        url = local_url or f"http://localhost:{server_args.webui_port}"
+        print(
+            f"""
+{delimiter}
+\033[1mSGLang Diffusion WebUI available at:\033[0m \033[1;4;92m{url}\033[0m
+{delimiter}
+"""
+        )
+
+        demo.block_thread()
--- a/python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py
+++ b/python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py
@@ -6,6 +6,7 @@ import argparse
 import os
 from typing import cast

+from sglang.multimodal_gen.apps.webui import run_sgl_diffusion_webui
 from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcommand
 from sglang.multimodal_gen.runtime.launch_server import launch_server
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
@@ -33,6 +34,9 @@ def execute_serve_cmd(args: argparse.Namespace, unknown_args: list[str] | None =
    server_args.post_init_serve()
    launch_server(server_args)

+    if server_args.webui:
+        run_sgl_diffusion_webui(server_args)
+

 class ServeSubcommand(CLISubcommand):
    """The `serve` subcommand for the sglang-diffusion CLI"""
--- a/python/sglang/multimodal_gen/runtime/launch_server.py
+++ b/python/sglang/multimodal_gen/runtime/launch_server.py
@@ -127,16 +127,28 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True):

    if launch_http_server:
        logger.info("Starting FastAPI server.")
+        if server_args.webui:
+            logger.info("Launch FastAPI server in another process because of webui.")
+            http_server_process = mp.Process(
+                target=launch_http_server_only,
+                args=(server_args,),
+                name=f"sglang-diffusion-webui",
+                daemon=True,
+            )
+            http_server_process.start()
+        else:
+            launch_http_server_only(server_args)

-        # set for endpoints to access global_server_args
-        set_global_server_args(server_args)

-        app = create_app(server_args)
-        uvicorn.run(
-            app,
-            use_colors=True,
-            log_level=server_args.log_level,
-            host=server_args.host,
-            port=server_args.port,
-            reload=False,
-        )
+def launch_http_server_only(server_args):
+    # set for endpoints to access global_server_args
+    set_global_server_args(server_args)
+    app = create_app(server_args)
+    uvicorn.run(
+        app,
+        log_config=None,
+        log_level=server_args.log_level,
+        host=server_args.host,
+        port=server_args.port,
+        reload=False,
+    )
--- a/python/sglang/multimodal_gen/runtime/server_args.py
+++ b/python/sglang/multimodal_gen/runtime/server_args.py
@@ -295,6 +295,10 @@ class ServerArgs:
    host: str | None = None
    port: int | None = None

+    # TODO: webui and their endpoint, check if webui_port is available.
+    webui: bool = False
+    webui_port: int | None = 12312
+
    scheduler_port: int = 5555

    # Stage verification
@@ -613,6 +617,19 @@ class ServerArgs:
            default=ServerArgs.port,
            help="Port for the HTTP API server.",
        )
+        parser.add_argument(
+            "--webui",
+            action=StoreBoolean,
+            default=ServerArgs.webui,
+            help="Whether to use webui for better display",
+        )
+
+        parser.add_argument(
+            "--webui-port",
+            type=int,
+            default=ServerArgs.webui_port,
+            help="Whether to use webui for better display",
+        )

        # Stage verification
        parser.add_argument(