showcase

app.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import base64
+import os
 import random
 import time
 import traceback
@@ -50,6 +51,7 @@ APP_DIR = Path(__file__).resolve().parent
 LOGO_PATH = APP_DIR / "assets" / "tryon_logo_compact.png"
 SHOWCASE_DIR = APP_DIR / "assets" / "showcase"
 SHOWCASE_IMAGES = [SHOWCASE_DIR / f"case_{idx}.png" for idx in range(4)]
+ZERO_GPU_DURATION_SECONDS = int(os.getenv("ZERO_GPU_DURATION_SECONDS", "240"))
 
 def _image_data_uri(path: Path) -> str:
     if not path.exists():
@@ -298,12 +300,12 @@ def _predict_from_prompt_files(prompt: str, files):
         )
 
 
-@spaces.GPU(size="xlarge", duration=180)
+@spaces.GPU(size="xlarge", duration=ZERO_GPU_DURATION_SECONDS)
 def predict_from_inputs(prompt: str, files):
     return _predict_from_prompt_files(prompt, files)
 
 
-@spaces.GPU(size="xlarge", duration=180)
+@spaces.GPU(size="xlarge", duration=ZERO_GPU_DURATION_SECONDS)
 def predict_from_message(message):
     prompt, files = _message_to_prompt_files(message)
     return _predict_from_prompt_files(prompt, files)
@@ -674,8 +676,11 @@ with gr.Blocks(title="JoyAI Virtual Try-On", css=CSS) as demo:
         inputs=[prompt_box],
         outputs=predict_outputs,
         api_name="tryon",
+        trigger_mode="once",
+        concurrency_limit=1,
+        concurrency_id="tryon_gpu",
     )
 
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(show_error=True)
+    demo.queue(max_size=4, default_concurrency_limit=1).launch(show_error=True)

inference.py

@@ -10,6 +10,7 @@ import random
 import re
 import sys
 import contextlib
+import threading
 import time
 from dataclasses import dataclass
 from pathlib import Path
@@ -73,9 +74,14 @@ class LoadedTryOnModel:
 
 _MODEL: LoadedTryOnModel | None = None
 _MODEL_DIR: Path | None = None
+_MODEL_LOCK = threading.Lock()
 _PE_CACHE: dict[str, str] = {}
 
 
+def _log_stage(message: str) -> None:
+    print(f"[tryon] {time.strftime('%Y-%m-%d %H:%M:%S')} | {message}", flush=True)
+
+
 def _get_bool_env(name: str, default: bool = False) -> bool:
     value = os.getenv(name)
     if value is None:
@@ -502,43 +508,61 @@ def load_model() -> LoadedTryOnModel:
     if _MODEL is not None:
         return _MODEL
 
-    _MODEL_DIR = _download_model_repo()
-    _add_code_paths(_MODEL_DIR)
-
-    from xvideo.config import load_config_class_from_pyfile
-    from xvideo.models import load_dit, load_pipeline
-
-    config_path = _find_default_config(_MODEL_DIR)
-    transformer_path = _find_transformer_path(_MODEL_DIR)
-    text_encoder_path = _resolve_text_encoder_path(_MODEL_DIR)
-    processor_path = _resolve_processor_path(_MODEL_DIR)
-    tokenizer_path = _resolve_tokenizer_path(_MODEL_DIR)
-
-    config_class = load_config_class_from_pyfile(str(config_path))
-    cfg = config_class()
-    cfg.dit_ckpt = str(transformer_path)
-    cfg.dit_ckpt_type = os.getenv(
-        "JOY_DIT_CKPT_TYPE",
-        "safetensor" if transformer_path.is_dir() else getattr(cfg, "dit_ckpt_type", "pt"),
-    )
-    cfg.training_mode = False
-    cfg.use_fsdp_inference = False
-    cfg.hsdp_shard_dim = 1
-    _patch_config_paths(cfg, _MODEL_DIR)
-
-    device = _device()
-    with _patch_split_firered_paths(
-        text_encoder_path=text_encoder_path,
-        processor_path=processor_path,
-        tokenizer_path=tokenizer_path,
-    ):
-        dit = load_dit(cfg, device=device)
-        dit.requires_grad_(False)
-        dit.eval()
-        pipeline = load_pipeline(cfg, dit, device)
-
-    _MODEL = LoadedTryOnModel(pipeline=pipeline, cfg=cfg, device=device, model_dir=_MODEL_DIR)
-    return _MODEL
+    with _MODEL_LOCK:
+        if _MODEL is not None:
+            return _MODEL
+
+        started_at = time.time()
+        _log_stage("load_model start")
+        _MODEL_DIR = _download_model_repo()
+        _add_code_paths(_MODEL_DIR)
+
+        from xvideo.config import load_config_class_from_pyfile
+        from xvideo.models import load_dit, load_pipeline
+
+        config_path = _find_default_config(_MODEL_DIR)
+        transformer_path = _find_transformer_path(_MODEL_DIR)
+        text_encoder_path = _resolve_text_encoder_path(_MODEL_DIR)
+        processor_path = _resolve_processor_path(_MODEL_DIR)
+        tokenizer_path = _resolve_tokenizer_path(_MODEL_DIR)
+        _log_stage(f"resolved config={config_path}")
+        _log_stage(f"resolved transformer={transformer_path}")
+        _log_stage(f"resolved text_encoder={text_encoder_path}")
+        _log_stage(f"resolved processor={processor_path}")
+        _log_stage(f"resolved tokenizer={tokenizer_path}")
+
+        config_class = load_config_class_from_pyfile(str(config_path))
+        cfg = config_class()
+        cfg.dit_ckpt = str(transformer_path)
+        cfg.dit_ckpt_type = os.getenv(
+            "JOY_DIT_CKPT_TYPE",
+            "safetensor" if transformer_path.is_dir() else getattr(cfg, "dit_ckpt_type", "pt"),
+        )
+        cfg.training_mode = False
+        cfg.use_fsdp_inference = False
+        cfg.hsdp_shard_dim = 1
+        _patch_config_paths(cfg, _MODEL_DIR)
+
+        device = _device()
+        _log_stage(f"loading DiT on {device}")
+        with _patch_split_firered_paths(
+            text_encoder_path=text_encoder_path,
+            processor_path=processor_path,
+            tokenizer_path=tokenizer_path,
+        ):
+            dit_started_at = time.time()
+            dit = load_dit(cfg, device=device)
+            _log_stage(f"loaded DiT in {time.time() - dit_started_at:.1f}s")
+            dit.requires_grad_(False)
+            dit.eval()
+            pipeline_started_at = time.time()
+            _log_stage("loading pipeline")
+            pipeline = load_pipeline(cfg, dit, device)
+            _log_stage(f"loaded pipeline in {time.time() - pipeline_started_at:.1f}s")
+
+        _MODEL = LoadedTryOnModel(pipeline=pipeline, cfg=cfg, device=device, model_dir=_MODEL_DIR)
+        _log_stage(f"load_model done in {time.time() - started_at:.1f}s")
+        return _MODEL
 
 
 def get_bucket_size(img: Image.Image) -> tuple[int, int]: