serving: manual load + GenerationMixin + use_cache=False + streaming

app.py

@@ -133,24 +133,28 @@ def stream(req: GenReq):
                                  media_type="text/plain; charset=utf-8", status_code=503)
     if (req.prompt or "").strip() == "":
         return StreamingResponse(iter(["(empty prompt)"]), media_type="text/plain")
-    if not _gen_lock.acquire(blocking=False):
-        return StreamingResponse(
-            iter(["The demo is busy generating for another visitor — please retry in a moment."]),
-            media_type="text/plain; charset=utf-8", status_code=429)
 
     from transformers import TextIteratorStreamer
 
     def run():
+        # Acquire AND release inside the generator. Acquiring in the endpoint and releasing in
+        # the generator's finally leaks the lock whenever the response is never consumed (client
+        # disconnects before streaming starts) — which wedges the demo at "busy" forever.
+        if not _gen_lock.acquire(timeout=2):
+            yield "The demo is busy generating for another visitor — please retry in a moment."
+            return
         try:
             tok, model, dev, ids, gen = _prep(req)
             streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True,
-                                            timeout=300)
-            th = threading.Thread(target=_safe_generate,
+                                            timeout=120)
+            th = threading.Thread(target=_safe_generate, daemon=True,
                                   args=(model, dict(input_ids=ids, streamer=streamer, **gen)))
             th.start()
             for chunk in streamer:
                 yield chunk
             th.join(timeout=5)
+        except Exception:
+            yield "\n[generation stopped]"
         finally:
             _gen_lock.release()