Reserve ZeroGPU time from the token budget; Gradio 6 theme/css on launch()

app.py

@@ -402,7 +402,27 @@ def _collect_references(
 # Inference
 # --------------------------------------------------------------------------- #
 
-@spaces.GPU(duration=150)
+DEFAULT_MAX_NEW_TOKENS = 1536
+
+# Measured on this Space: ~34 decoded tokens/s plus ~2.5 s of encode/prefill.
+# Reserving from the token budget keeps the ZeroGPU hold tight instead of
+# padding every visitor's quota with a fixed worst case.
+MEASURED_TOKENS_PER_SECOND = 32.0
+
+
+def _gpu_duration(*args, **kwargs) -> int:
+    """Reserve ZeroGPU time from the requested token budget."""
+    tokens = kwargs.get("max_new_tokens")
+    if tokens is None and len(args) >= 14:
+        tokens = args[13]
+    try:
+        tokens = int(tokens)
+    except (TypeError, ValueError):
+        tokens = DEFAULT_MAX_NEW_TOKENS
+    return int(min(150, max(25, round(6 + tokens / MEASURED_TOKENS_PER_SECOND))))
+
+
+@spaces.GPU(duration=_gpu_duration)
 def rewrite(
     task: str = "T2AV",
     prompt: str = "",
@@ -417,7 +437,7 @@ def rewrite(
     greedy: bool = True,
     temperature: float = 0.7,
     top_p: float = 0.9,
-    max_new_tokens: int = 2048,
+    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
     seed: int = 42,
     video_fps: float = 1.0,
 ):
@@ -624,7 +644,7 @@ def _on_task_change(task: str):
     )
 
 
-with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="MiniMax-H3 Prompt Rewriter") as demo:
+with gr.Blocks(title="MiniMax-H3 Prompt Rewriter") as demo:
     gr.Markdown(
         "# MiniMax-H3 Prompt Rewriter · Qwen2.5-Omni LoRA\n"
         "Turn a short request — plus optional image, video or audio references — into a "
@@ -703,8 +723,9 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="MiniMax-H3 Prompt Rewri
                 )
             with gr.Row():
                 max_new_tokens = gr.Slider(
-                    minimum=256, maximum=4096, value=2048, step=128,
+                    minimum=256, maximum=4096, value=DEFAULT_MAX_NEW_TOKENS, step=128,
                     label="Max new tokens",
+                    info="Also sets the reserved ZeroGPU time (~32 tokens/s).",
                 )
                 seed = gr.Number(value=42, precision=0, label="Seed")
                 video_fps = gr.Slider(
@@ -802,10 +823,11 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="MiniMax-H3 Prompt Rewri
             image_1, image_2, image_3, image_4, ref_video, ref_audio,
             resolution, task_help,
         ],
+        api_name=False,
     )
 
     run_btn.click(fn=rewrite, inputs=ALL_INPUTS, outputs=[output, status], api_name="rewrite")
     prompt.submit(fn=rewrite, inputs=ALL_INPUTS, outputs=[output, status], api_name=False)
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)