Update app.py

app.py

@@ -68,14 +68,14 @@ processor = AutoProcessor.from_pretrained(MODEL_ID)
 
 
 # -----------------------------
-# Streaming generator
+# Inference (single-shot)
 # -----------------------------
 @spaces.GPU
-def analyze_ecg_stream(image: Image.Image):
+def analyze_ecg(image: Image.Image):
     if image is None:
-        yield "Please upload an ECG image."
-        return
+        return "Please upload an ECG image."
 
+    # Chat messages (preferred) -> insert image via template
     messages = [
         {"role": "user", "content": [
             {"type": "text", "text": CLINICAL_PROMPT},
@@ -83,95 +83,82 @@ def analyze_ecg_stream(image: Image.Image):
         ]}
     ]
 
-    # Prepare inputs (try chat template; fallback to plain text+image)
     try:
         chat_text = processor.apply_chat_template(messages, add_generation_prompt=True)
-        model_inputs = processor(text=chat_text, images=image, return_tensors="pt")
+        inputs = processor(text=chat_text, images=image, return_tensors="pt")
     except Exception:
-        # Fallback when the template/image token count mismatches
-        model_inputs = processor(text=CLINICAL_PROMPT, images=image, return_tensors="pt")
-        yield "[Note] Using fallback prompt packing.\n"
+        # Fallback if image-token count/template mismatches
+        inputs = processor(text=CLINICAL_PROMPT, images=image, return_tensors="pt")
 
-    model_inputs = {k: v.to(DEVICE) for k, v in model_inputs.items()}
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
 
-    # Streamer must use the tokenizer (not the processor)
-    streamer = TextIteratorStreamer(
-        processor.tokenizer, skip_prompt=True, skip_special_tokens=True
-    )
-
-    def _generate():
-        try:
-            model.generate(
-                **model_inputs,
-                streamer=streamer,
-                **GEN_KW,
-            )
-        except Exception as e:
-            streamer.put("\n\n[Generation Error]\n" + traceback.format_exc())
-        finally:
-            streamer.end()
-
-    t = threading.Thread(target=_generate, daemon=True)
-    t.start()
-
-    buf = ""
-    for piece in streamer:
-        buf += piece
-        yield buf
+    try:
+        with torch.inference_mode():
+            out = model.generate(**inputs, **GEN_KW)
+
+        # Decode only new tokens if possible
+        input_len = inputs.get("input_ids", None)
+        if input_len is not None:
+            in_len = inputs["input_ids"].shape[-1]
+            gen_ids = out[0][in_len:] if out.shape[-1] > in_len else out[0]
+        else:
+            gen_ids = out[0]
+
+        text = processor.decode(gen_ids, skip_special_tokens=True).strip()
+        return text if text else "[Empty output] Try a clearer ECG image."
+    except Exception as e:
+        return f"[Generation Error]\n{traceback.format_exc()}"
 
 def reset():
     return None, ""
 
 # -----------------------------
-# UI
+# UI (accessible contrast)
 # -----------------------------
 theme = gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")
+css = """
+/* High-contrast output area */
+textarea[aria-label="Generated ECG Report"] {
+  background-color: #ffffff !important;
+  color: #111827 !important;   /* near-black */
+}
+
+/* Card look */
+.card { background:#fff; border:1px solid #e5e7eb; border-radius:14px; padding:16px; }
 
-custom_css = """
-#app {
-  max-width: 1100px;
-  margin: 0 auto;
+/* Buttons */
+.gr-button { background-color:#1e3a8a !important; color:#ffffff !important; }
+
+/* Disclaimer with strong contrast */
+.disclaimer {
+  margin-top:12px; padding:12px 14px; border-radius:12px;
+  background:#fff7ed; color:#7c2d12; border:1px solid #fdba74; font-weight:600;
 }
+
+/* Header */
 .header {
   display:flex; align-items:center; justify-content:space-between;
-  padding: 16px 14px; border-radius: 14px;
-  background: linear-gradient(135deg, #1f2937 0%, #111827 100%);
-  color: #fff; box-shadow: 0 6px 20px rgba(0,0,0,0.25);
-}
-.brand { font-size: 18px; font-weight: 700; letter-spacing: 0.3px; }
-.disclaimer {
-  margin-top: 12px; padding: 12px 14px; border-radius: 12px;
-  background: #fef2f2; color:#7f1d1d; border:1px solid #fecaca; font-weight:600;
+  padding:14px 16px; border-radius:14px; background:#0f172a; color:#e5e7eb;
 }
-.card {
-  background: #ffffff; border: 1px solid #e5e7eb; border-radius: 14px;
-  padding: 16px; box-shadow: 0 8px 18px rgba(17,24,39,0.06);
-}
-footer {
-  font-size: 12px; color:#6b7280; margin-top: 8px;
-}
-.gr-button { background-color:#1e3a8a !important; color:#fff !important; }
+.header .brand { font-weight:800; letter-spacing:.2px; }
 """
 
-with gr.Blocks(theme=theme, css=custom_css, elem_id="app") as demo:
-    with gr.Row():
-        gr.HTML("""
-        <div class="header">
-          <div class="brand">🩺 ECG Interpretation Assistant</div>
-          <div>Gemma-ECG-Vision</div>
-        </div>
-        <div class="disclaimer">
-          ⚠️ <strong>Education & Research Only:</strong> This tool is not a medical device and must not be used for diagnosis or treatment.
-          Always consult a licensed clinician for interpretation and clinical decisions.
-        </div>
-        """)
+with gr.Blocks(theme=theme, css=css) as demo:
+    gr.HTML(f"""
+    <div class="header">
+      <div class="brand">🫀 ECG Interpretation Assistant</div>
+      <div>Model: <code>{MODEL_ID}</code></div>
+    </div>
+    <div class="disclaimer">
+      ⚠️ <strong>Education & Research Only:</strong> This tool is not a medical device and must not be used for diagnosis or treatment.
+      Always consult a licensed clinician for interpretation and clinical decisions.
+    </div>
+    """)
 
     with gr.Row(equal_height=True):
         with gr.Column(scale=1):
             with gr.Group(elem_classes="card"):
-                image_input = gr.Image(
-                    type="pil", label="Upload ECG Image", height=360, show_label=True
-                )
+                image_input = gr.Image(type="pil", label="Upload ECG Image", height=360)
                 with gr.Row():
                     submit_btn = gr.Button("Generate Report", variant="primary")
                     reset_btn = gr.Button("Reset")
@@ -179,43 +166,17 @@ with gr.Blocks(theme=theme, css=custom_css, elem_id="app") as demo:
         with gr.Column(scale=2):
             with gr.Group(elem_classes="card"):
                 output_box = gr.Textbox(
-                    label="Generated ECG Report (Streaming)",
-                    lines=26,
+                    label="Generated ECG Report",
+                    lines=28,
                     show_copy_button=True,
-                    autoscroll=True,
                     placeholder="The model's report will appear here…",
                 )
-                gr.Markdown(
-                    "Tip: Clear, high-resolution ECGs with visible lead labels improve P wave and ST-segment assessment."
-                )
 
-    gr.HTML(f"""
-    <footer>
-      Model: <code>{MODEL_ID}</code> | Device: <code>{DEVICE}</code>
-    </footer>
-    """)
-
-    submit_btn.click(
-        fn=analyze_ecg_stream,
-        inputs=image_input,
-        outputs=output_box,
-        queue=True,
-        api_name="analyze_ecg",
+    gr.Markdown(
+        "Tip: Use a high-resolution ECG with visible lead labels to improve assessment of P waves and ST segments."
     )
+
+    submit_btn.click(analyze_ecg, inputs=image_input, outputs=output_box, queue=False)
     reset_btn.click(reset, outputs=[image_input, output_box])
 
-def queue_with_compat(demo, max_size=32, limit=4):
-    params = inspect.signature(gr.Blocks.queue).parameters
-    if "concurrency_count" in params:
-        # Older Gradio 3.x / early 4.x
-        return demo.queue(concurrency_count=limit, max_size=max_size)
-    elif "default_concurrency_limit" in params:
-        # Newer Gradio 4.x
-        return demo.queue(default_concurrency_limit=limit, max_size=max_size)
-    else:
-        # Fallback – no knobs exposed
-        return demo.queue()
-
-# … build your UI as before …
-queue_with_compat(demo, max_size=32, limit=4)
 demo.launch(share=False, debug=True)