update app

app.py

@@ -19,6 +19,7 @@ from transformers import (
     Qwen2_5_VLForConditionalGeneration,
     Qwen3_5ForConditionalGeneration,
     Qwen3VLForConditionalGeneration,
+    Gemma4ForConditionalGeneration,
     AutoProcessor,
     AutoModelForImageTextToText,
     TextIteratorStreamer,
@@ -42,6 +43,7 @@ LFM_450_MODEL_NAME   = "LiquidAI/LFM2.5-VL-450M"
 LFM_16_MODEL_NAME    = "LiquidAI/LFM2.5-VL-1.6B"
 QWEN_UNREDACTED_NAME = "prithivMLmods/Qwen3.5-2B-Unredacted-MAX"
 QWEN25_VL_3B_NAME    = "Qwen/Qwen2.5-VL-3B-Instruct"
+GEMMA4_E2B_NAME      = "google/gemma-4-E2B-it"
 
 # ── Qwen3.5-4B ──────────────────────────────────────────
 print(f"Loading Qwen3.5-4B model: {QWEN_4B_MODEL_NAME} on {DEVICE}...")
@@ -144,6 +146,23 @@ except Exception as e:
     qwen25_vl_3b_model = None
     qwen25_vl_3b_processor = None
 
+# ── Gemma4-E2B-it ───────────────────────────────────────
+print(f"Loading Gemma4-E2B-it: {GEMMA4_E2B_NAME} on {DEVICE}...")
+try:
+    gemma4_e2b_model = Gemma4ForConditionalGeneration.from_pretrained(
+        GEMMA4_E2B_NAME,
+        torch_dtype=torch.bfloat16,
+        device_map="auto" if torch.cuda.is_available() else None,
+    ).eval()
+    if not torch.cuda.is_available():
+        gemma4_e2b_model = gemma4_e2b_model.to(DEVICE)
+    gemma4_e2b_processor = AutoProcessor.from_pretrained(GEMMA4_E2B_NAME)
+    print("Gemma4-E2B-it model loaded successfully.")
+except Exception as e:
+    print(f"Warning: Gemma4-E2B-it model loading failed. Error: {e}")
+    gemma4_e2b_model = None
+    gemma4_e2b_processor = None
+
 
 # --- Utility Functions ---
 def safe_parse_json(text: str):
@@ -393,6 +412,40 @@ def generate_inference_stream(
                 yield f"data: {json.dumps({'chunk': tok})}\n\n"
         thread.join()
 
+    # ── Gemma4-E2B-it ───────────────────────────────────
+    elif model_id == "gemma4_e2b":
+        if gemma4_e2b_model is None or gemma4_e2b_processor is None:
+            yield f"data: {json.dumps({'chunk': '[Error] Gemma4-E2B-it model not loaded.'})}\n\n"
+            yield "data: [DONE]\n\n"
+            return
+        messages = [{"role": "user", "content": [
+            {"type": "image", "image": image},
+            {"type": "text",  "text":  full_prompt},
+        ]}]
+        text_input = gemma4_e2b_processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        inputs = gemma4_e2b_processor(
+            text=[text_input],
+            images=[image],
+            return_tensors="pt",
+            padding=True,
+        ).to(gemma4_e2b_model.device)
+        streamer = TextIteratorStreamer(
+            gemma4_e2b_processor.tokenizer,
+            skip_prompt=True, skip_special_tokens=True, timeout=120,
+        )
+        thread = threading.Thread(
+            target=gemma4_e2b_model.generate,
+            kwargs=dict(**inputs, streamer=streamer, max_new_tokens=1024,
+                        use_cache=True, temperature=1.0, do_sample=True),
+        )
+        thread.start()
+        for tok in streamer:
+            if tok:
+                yield f"data: {json.dumps({'chunk': tok})}\n\n"
+        thread.join()
+
     yield "data: [DONE]\n\n"
 
 
@@ -745,6 +798,7 @@ async def homepage(request: Request):
         .model-badge.lfm16   { background: rgba(107,203,119,0.15); color: #6bcb77;        border: 1px solid rgba(107,203,119,0.35); }
         .model-badge.qunred  { background: rgba(255,80,160,0.15);  color: #ff50a0;        border: 1px solid rgba(255,80,160,0.35); }
         .model-badge.q25vl3b { background: rgba(80,180,255,0.15);  color: #50b4ff;        border: 1px solid rgba(80,180,255,0.35); }
+        .model-badge.g4e2b   { background: rgba(66,197,107,0.15);  color: #42c56b;        border: 1px solid rgba(66,197,107,0.35); }
 
         .model-info-box {
             border-radius: 6px; padding: 9px;
@@ -761,7 +815,7 @@ async def homepage(request: Request):
     <span class="logo">MULTIMODAL EDGE</span>
     <span class="sep">|</span>
     <span class="sub">Node-Based Inference Canvas</span>
-    <span class="badge">v2.5 — HEPTA MODEL</span>
+    <span class="badge">v2.6 — OCTA MODEL</span>
 </div>
 
 <div id="canvas">
@@ -836,6 +890,7 @@ async def homepage(request: Request):
                     <option value="lfm_16">LFM2.5-VL-1.6B (LiquidAI)</option>
                     <option value="qwen_unredacted">Qwen3.5-2B-Unredacted-MAX</option>
                     <option value="qwen25_vl_3b">Qwen2.5-VL-3B-Instruct</option>
+                    <option value="gemma4_e2b">Gemma4-E2B-it (Google)</option>
                 </select>
             </div>
             <div id="modelInfoBox" class="model-info-box"
@@ -1091,6 +1146,13 @@ const MODEL_INFO = {
         bg:     'rgba(80,180,255,0.07)',
         border: 'rgba(80,180,255,0.25)',
     },
+    gemma4_e2b: {
+        html:   `<span class="model-badge g4e2b">GEMMA 4 · E2B</span><br><br>
+                 Gemma4-E2B-it by Google DeepMind. Efficient 2B multimodal model
+                 with strong vision-language understanding &amp; instruction-following.`,
+        bg:     'rgba(66,197,107,0.07)',
+        border: 'rgba(66,197,107,0.25)',
+    },
 };
 
 modelSelect.onchange = () => {