Build the model once in the worker, size the GPU reservation by the length, and check requests before the GPU is reserved.

app.py

@@ -44,9 +44,15 @@ def _unique_path(ext: str) -> str:
     return os.path.join("/tmp", f"t2m_{time_ns()}_{uuid4().hex[:8]}{ext}")
 
 
-# Core Text -> MIDI
-def generate_midi(prompt: str, temperature: float = 0.9, max_len: int = 500) -> str:
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+model = None
+r_tokenizer = None
+tokenizer = None
+
+
+def load_model():
+    """Builds the model on the CPU, then places it on DEVICE."""
+    global model, r_tokenizer, tokenizer
 
     # Load REMI vocab/tokenizer (pickle dict used by the provided model)
     with open(TOKENIZER_PATH, "rb") as f:
@@ -62,16 +68,28 @@ def generate_midi(prompt: str, temperature: float = 0.9, max_len: int = 500) ->
         1024,          # max_seq_len
         False,         # use_rotary
         8,             # rotary_dim
-        device=device  # device
+        device="cpu"   # device
     )
-    model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
+    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
     model.eval()
+    model.to(DEVICE)
 
     tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
+
+
+# Built in the worker, which runs process_fn, on its main thread, where ZeroGPU's emulation records
+# the CUDA placement, so that no request spends the caller's allowance on building it. The server
+# never runs process_fn.
+if __name__ != "__main__":
+    load_model()
+
+
+# Core Text -> MIDI
+def generate_midi(prompt: str, temperature: float = 0.9, max_len: int = 500) -> str:
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
 
-    input_ids = nn.utils.rnn.pad_sequence(inputs.input_ids, batch_first=True, padding_value=0).to(device)
-    attention_mask = nn.utils.rnn.pad_sequence(inputs.attention_mask, batch_first=True, padding_value=0).to(device)
+    input_ids = nn.utils.rnn.pad_sequence(inputs.input_ids, batch_first=True, padding_value=0).to(DEVICE)
+    attention_mask = nn.utils.rnn.pad_sequence(inputs.attention_mask, batch_first=True, padding_value=0).to(DEVICE)
 
     with torch.no_grad():
         output = model.generate(input_ids, attention_mask, max_len=max_len, temperature=temperature)
@@ -83,7 +101,10 @@ def generate_midi(prompt: str, temperature: float = 0.9, max_len: int = 500) ->
     generated_midi.dump_midi(midi_path)
     return midi_path
 
-@spaces.GPU(duration=120)
+def gpu_seconds(prompt: str, temperature: float, max_length: int) -> float:
+    return 15 + 0.1 * int(max_length)
+
+
 def process_fn(prompt: str, temperature: float, max_length: int):
     """
     Generate a MIDI from a text description.
@@ -98,6 +119,15 @@ def process_fn(prompt: str, temperature: float, max_length: int):
         midi_path (str): The generated MIDI.
     """
 
+    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
+    # caller's ZeroGPU allowance
+    if not (prompt or "").strip():
+        raise gr.Error("Describe the music to generate.")
+    return compose(prompt, temperature, max_length)
+
+
+@spaces.GPU(duration=gpu_seconds)
+def compose(prompt: str, temperature: float, max_length: int) -> str:
     try:
         return generate_midi(prompt, float(temperature), int(max_length))
     except Exception as error: