Update app.py

app.py

@@ -36,10 +36,8 @@ MODEL_MAP = {
     "GPT-3.5 (gpt-3.5-turbo) - OpenAI": {"kind": "openai-chat", "id": "gpt-3.5-turbo"},
 }
 
-# Cache for loaded Hugging Face models/pipelines
 HF_PIPELINES = {}
 
-# OpenAI client (only if key exists)
 OPENAI_KEY = os.getenv("OPENAI_API_KEY")
 OPENAI_CLIENT = OpenAI(api_key=OPENAI_KEY) if OPENAI_KEY else None
 
@@ -55,19 +53,14 @@ def get_hf_pipeline(model_id: str):
     mdl = AutoModelForCausalLM.from_pretrained(
         model_id,
         low_cpu_mem_usage=True,
-        torch_dtype=torch.float32,  # CPU-safe
+        torch_dtype=torch.float32,
     )
 
     # Some older models (e.g., GPT-1/2) have no pad token
     if tok.pad_token_id is None and tok.eos_token_id is not None:
         tok.pad_token = tok.eos_token
 
-    gen = pipeline(
-        "text-generation",
-        model=mdl,
-        tokenizer=tok,
-        device=device,
-    )
+    gen = pipeline("text-generation", model=mdl, tokenizer=tok, device=device)
     HF_PIPELINES[model_id] = gen
     return gen
 
@@ -92,9 +85,7 @@ def generate_stream(model_choice, prompt, max_new_tokens, temperature, top_p, se
             tok = gen.tokenizer
             mdl = gen.model
 
-            streamer = TextIteratorStreamer(
-                tok, skip_prompt=True, skip_special_tokens=True
-            )
+            streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
 
             inputs = tok(prompt, return_tensors="pt")
             if torch.cuda.is_available():
@@ -111,7 +102,6 @@ def generate_stream(model_choice, prompt, max_new_tokens, temperature, top_p, se
                 streamer=streamer,
             )
 
-            # Run generation in a thread so we can iterate streamer
             thread = Thread(target=mdl.generate, kwargs=generate_kwargs)
             thread.start()
 
@@ -191,7 +181,6 @@ with gr.Blocks(title="Mini GPT Playground") as demo:
 
     model_choice.change(maybe_warn, inputs=[model_choice], outputs=[warn])
 
-    # Streamed generation
     generate_btn.click(
         fn=generate_stream,
         inputs=[model_choice, prompt, max_new_tokens, temperature, top_p, seed],
@@ -199,30 +188,26 @@ with gr.Blocks(title="Mini GPT Playground") as demo:
     )
 
 # -------------------------
-# Robust initialization for HF Spaces
+# Spaces-friendly init
 # -------------------------
-# 1) Try to mount into a FastAPI app (works for "Python" Spaces)
-# 2) Otherwise, fall back to launching Gradio directly
+# 1) Expose FastAPI app if running in a Python SDK Space (optional)
 app = None
-if os.getenv("SPACE_ID"):
+try:
+    # If FastAPI is available, we provide an app so Python SDK Spaces can import it.
+    from fastapi import FastAPI
+    app = FastAPI()
     try:
-        from fastapi import FastAPI
-        app = FastAPI()
-        # Optional: small queue (if available in your Gradio version)
-        try:
-            demo = demo.queue(max_size=8)
-        except TypeError:
-            pass
-        app = gr.mount_gradio_app(app, demo, path="/")
-    except Exception:
-        # FastAPI not available or mount failed; we'll rely on launch below when run locally
-        app = None
-
-# For local dev or plain Gradio Spaces
+        demo = demo.queue(max_size=8)
+    except TypeError:
+        pass
+    app = gr.mount_gradio_app(app, demo, path="/")
+except Exception:
+    app = None  # fine on Gradio SDK Spaces
+
+# 2) For local runs / Gradio SDK Spaces: DO NOT set server_port; let Gradio pick the env port.
 if __name__ == "__main__":
-    # Version-safe queue/launch
     try:
         demo = demo.queue(max_size=8)
     except TypeError:
         pass
-    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))
+    demo.launch()  # no server_port — avoids port collisions on Spaces