Update app.py

app.py

@@ -3,15 +3,13 @@ import gradio as gr
 import subprocess, os, sys, time
 import urllib.request, urllib.error
 from fastapi import Request
-from fastapi.responses import Response, PlainTextResponse, HTMLResponse, StreamingResponse
+from fastapi.responses import Response, PlainTextResponse, HTMLResponse
 
 PORT = 9001
 FM_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "file-manager")
 MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 
-# =================================================================
-# Launch the file manager as a background child process
-# =================================================================
+# --- Launch the file manager as a background child process ---
 proc = {"p": None}
 def start_fm():
     if proc["p"] and proc["p"].poll() is None:
@@ -25,9 +23,7 @@ def start_fm():
 
 start_fm()
 
-# =================================================================
-# GPU info: runs nvidia-smi INSIDE the gate (GPU only attached here)
-# =================================================================
+# --- GPU info (runs INSIDE the gate) ---
 @spaces.GPU
 def gpu_report():
     import torch
@@ -45,9 +41,7 @@ def gpu_report():
         lines.append(f"nvidia-smi error: {e}")
     return "\n".join(lines)
 
-# =================================================================
-# LLM inference: model loads + generates INSIDE the gate
-# =================================================================
+# --- LLM (loads + generates INSIDE the gate) ---
 _state = {"model": None, "tok": None}
 
 @spaces.GPU
@@ -66,11 +60,9 @@ def run_llm(prompt: str):
     reply = tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     return f"[GPU: {torch.cuda.get_device_name(0)}]\n\n{reply}"
 
-# =================================================================
-# Gradio owns the ONLY public port 7860
-# =================================================================
+# --- Gradio owns the ONLY public port 7860 ---
 with gr.Blocks() as demo:
-    gr.Markdown("Self-editing computer. File manager + terminal + Python REPL at /fm")
+    gr.Markdown("Self-editing computer. Terminal + file manager at /fm")
     with gr.Row():
         prompt = gr.Textbox(label="Ask the LLM")
         out = gr.Textbox(label="Reply")
@@ -82,36 +74,24 @@ demo.queue()
 demo.launch(server_name="0.0.0.0", server_port=7860,
             ssr_mode=False, prevent_thread_lock=True)
 
-app = demo.app  # Gradio's FastAPI app on 7860
+app = demo.app
 
-# =================================================================
-# Hot-reload: respawn ONLY the file-manager process, no Space rebuild
-# =================================================================
+# --- Hot-reload: respawn ONLY the file-manager process, no Space rebuild ---
 @app.get("/reload")
 def reload_fm():
     start_fm()
     time.sleep(1)
     return HTMLResponse('Reloaded file-manager process. <a href="/fm/">back to /fm</a>')
 
-# =================================================================
-# GPU info route
-# =================================================================
 @app.get("/gpuinfo")
 def gpuinfo():
     return PlainTextResponse(gpu_report())
 
-# =================================================================
-# LLM route
-# =================================================================
 @app.get("/llm")
 def llm(q: str = "Say hello and tell me which GPU you are running on."):
     return PlainTextResponse(run_llm(q))
 
-# =================================================================
-# Proxy: /fm/... -> 127.0.0.1:9001/fm/...
-# Streams text/plain (live shell) chunk by chunk; buffers everything else.
-# Forwards method + body so POST forms and the JSON REPL work.
-# =================================================================
+# --- Proxy: /fm/... -> 127.0.0.1:9001/fm/... (method + body forwarded) ---
 @app.api_route("/fm", methods=["GET", "POST"])
 @app.api_route("/fm/{path:path}", methods=["GET", "POST"])
 async def proxy(request: Request, path: str = ""):
@@ -122,33 +102,15 @@ async def proxy(request: Request, path: str = ""):
     fwd = {k: v for k, v in request.headers.items() if k.lower() != "host"}
     req = urllib.request.Request(url, data=body or None, method=request.method, headers=fwd)
     try:
-        r = urllib.request.urlopen(req)
+        with urllib.request.urlopen(req) as r:
+            data, status, headers = r.read(), r.status, r.headers
     except urllib.error.HTTPError as e:
-        return Response(
-            content=e.read(), status_code=e.code,
-            headers={"content-type": e.headers.get("content-type", "text/plain")})
+        data, status, headers = e.read(), e.code, e.headers
     except urllib.error.URLError:
         return Response(content=b"file-manager not up yet on 9001", status_code=502)
+    out = {k: v for k, v in headers.items()
+           if k.lower() in ("set-cookie", "location", "content-type")}
+    return Response(content=data, status_code=status, headers=out)
 
-    ctype = r.headers.get("content-type", "text/html")
-    # Stream the live shell output; buffer normal pages
-    if ctype.startswith("text/plain"):
-        def gen():
-            while True:
-                chunk = r.read(1024)
-                if not chunk:
-                    break
-                yield chunk
-        return StreamingResponse(gen(), media_type="text/plain")
-
-    out_headers = {"content-type": ctype}
-    for k, v in r.headers.items():
-        if k.lower() in ("set-cookie", "location"):
-            out_headers[k] = v
-    return Response(content=r.read(), status_code=r.status, headers=out_headers)
-
-# =================================================================
-# Keep the process alive (launch returned control)
-# =================================================================
 while True:
     time.sleep(3600)