fix app.py

app.py

@@ -1,6 +1,7 @@
 import argparse
 import functools
 import importlib.util
+import os
 from pathlib import Path
 import re
 import time
@@ -11,6 +12,19 @@ import numpy as np
 import torch
 from transformers import AutoModel, AutoProcessor
 
+try:
+    import spaces
+except ImportError:
+    class _SpacesFallback:
+        @staticmethod
+        def GPU(*_args, **_kwargs):
+            def _decorator(func):
+                return func
+
+            return _decorator
+
+    spaces = _SpacesFallback()
+
 # Disable the broken cuDNN SDPA backend
 torch.backends.cuda.enable_cudnn_sdp(False)
 # Keep these enabled as fallbacks
@@ -21,6 +35,7 @@ torch.backends.cuda.enable_math_sdp(True)
 MODEL_PATH = "OpenMOSS-Team/MOSS-TTS"
 DEFAULT_ATTN_IMPLEMENTATION = "auto"
 DEFAULT_MAX_NEW_TOKENS = 4096
+PRELOAD_ENV_VAR = "MOSS_TTS_PRELOAD_AT_STARTUP"
 CONTINUATION_NOTICE = (
     "Continuation mode is active. Make sure the reference audio transcript is prepended to the input text."
 )
@@ -289,6 +304,7 @@ def apply_example_selection(
     )
 
 
+@spaces.GPU(duration=180)
 def run_inference(
     text: str,
     reference_audio: str | None,
@@ -574,48 +590,102 @@ def build_demo(args: argparse.Namespace):
     return demo
 
 
+def resolve_runtime_attn(args: argparse.Namespace) -> argparse.Namespace:
+    runtime_device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+    runtime_dtype = torch.bfloat16 if runtime_device.type == "cuda" else torch.float32
+    args.attn_implementation = resolve_attn_implementation(
+        requested=args.attn_implementation,
+        device=runtime_device,
+        dtype=runtime_dtype,
+    ) or "none"
+    return args
+
+
+def parse_bool_env(name: str, default: bool) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "y", "on"}
+
+
+def parse_port(value: str | None, default: int) -> int:
+    if not value:
+        return default
+    try:
+        return int(value)
+    except ValueError:
+        return default
+
+
+def build_default_args() -> argparse.Namespace:
+    return resolve_runtime_attn(
+        argparse.Namespace(
+            model_path=MODEL_PATH,
+            device="cuda:0",
+            attn_implementation=DEFAULT_ATTN_IMPLEMENTATION,
+            host=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"),
+            port=parse_port(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT")), 7860),
+            share=False,
+        )
+    )
+
+
 def main():
     parser = argparse.ArgumentParser(description="MossTTS Gradio Demo")
     parser.add_argument("--model_path", type=str, default=MODEL_PATH)
     parser.add_argument("--device", type=str, default="cuda:0")
     parser.add_argument("--attn_implementation", type=str, default=DEFAULT_ATTN_IMPLEMENTATION)
     parser.add_argument("--host", type=str, default="0.0.0.0")
-    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=int(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT", "7860"))),
+    )
     parser.add_argument("--share", action="store_true")
     args = parser.parse_args()
 
-    runtime_device = torch.device(args.device if torch.cuda.is_available() else "cpu")
-    runtime_dtype = torch.bfloat16 if runtime_device.type == "cuda" else torch.float32
-    args.attn_implementation = resolve_attn_implementation(
-        requested=args.attn_implementation,
-        device=runtime_device,
-        dtype=runtime_dtype,
-    ) or "none"
+    args.host = os.getenv("GRADIO_SERVER_NAME", args.host)
+    args.port = parse_port(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT")), args.port)
+    args = resolve_runtime_attn(args)
     print(f"[INFO] Using attn_implementation={args.attn_implementation}", flush=True)
 
-    # Preload model/processor at startup to avoid first-request cold start latency.
-    preload_started_at = time.monotonic()
-    print(
-        f"[Startup] Preloading backend: model={args.model_path}, device={args.device}, attn={args.attn_implementation}",
-        flush=True,
-    )
-    load_backend(
-        model_path=args.model_path,
-        device_str=args.device,
-        attn_implementation=args.attn_implementation,
-    )
-    print(
-        f"[Startup] Backend preload finished in {time.monotonic() - preload_started_at:.2f}s",
-        flush=True,
-    )
+    preload_enabled = parse_bool_env(PRELOAD_ENV_VAR, default=not bool(os.getenv("SPACE_ID")))
+    if preload_enabled:
+        preload_started_at = time.monotonic()
+        print(
+            f"[Startup] Preloading backend: model={args.model_path}, device={args.device}, attn={args.attn_implementation}",
+            flush=True,
+        )
+        load_backend(
+            model_path=args.model_path,
+            device_str=args.device,
+            attn_implementation=args.attn_implementation,
+        )
+        print(
+            f"[Startup] Backend preload finished in {time.monotonic() - preload_started_at:.2f}s",
+            flush=True,
+        )
+    else:
+        print(
+            f"[Startup] Skipping preload (set {PRELOAD_ENV_VAR}=1 to enable).",
+            flush=True,
+        )
 
     demo = build_demo(args)
     demo.queue(max_size=16, default_concurrency_limit=1).launch(
         server_name=args.host,
         server_port=args.port,
         share=args.share,
+        ssr_mode=False,
     )
 
 
+# Expose a module-level demo for Gradio hot-reload/Spaces launcher.
+demo = build_demo(build_default_args())
+
+
 if __name__ == "__main__":
-    main()
+    if os.getenv("GRADIO_HOT_RELOAD"):
+        print("[Startup] GRADIO_HOT_RELOAD detected. Skipping explicit launch().", flush=True)
+    else:
+        main()