Rename app.py to main.py

app.py → main.py

@@ -1,6 +1,5 @@
 import os
 import torch
-from threading import Thread
 from PIL import Image
 import numpy as np
 
@@ -11,18 +10,16 @@ import gradio as gr
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
-    TextIteratorStreamer,
     AutoProcessor,
 )
 
-# Force CPU
+# Force CPU Only
 os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
 torch.cuda.is_available = lambda: False
 device = "cpu"
 print("Running on CPU ✅")
 
-# ---------------- LOAD MAIN CHAT MODEL ----------------
-
+# ---------------- LOAD CHAT MODEL ----------------
 MODEL_ID = "microsoft/Phi-3.5-mini-instruct"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -33,8 +30,7 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 ).eval()
 
-# ---------------- LOAD VISION MODEL (FlashAttention disabled) ----------------
-
+# ---------------- LOAD VISION MODEL ----------------
 models = {}
 processors = {}
 
@@ -46,48 +42,32 @@ try:
         torch_dtype=torch.float32,
         device_map="cpu",
         low_cpu_mem_usage=True,
-        attn_implementation="eager"   # <<< KEY FIX ✅
+        attn_implementation="eager"
     ).eval()
 
     processors[VISION_ID] = AutoProcessor.from_pretrained(
         VISION_ID,
         trust_remote_code=True
     )
-
     print("Vision model loaded ✅")
 except Exception as e:
     print("Vision model failed to load:", e)
 
-# ---------------- CHAT FUNCTION (for UI) ----------------
-
+# ---------------- CHAT FUNCTION ----------------
 def chat_simple(message, history):
-    system_prompt = "You are a helpful assistant."
-    
-    conversation = [{"role": "system", "content": system_prompt}]
+    conversation = [{"role": "system", "content": "You are a helpful assistant."}]
     for user, assistant in history:
         conversation.append({"role": "user", "content": user})
         conversation.append({"role": "assistant", "content": assistant})
     conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(
-        conversation,
-        add_generation_prompt=True,
-        return_tensors="pt"
-    ).to(device)
-
-    output = model.generate(
-        input_ids,
-        max_new_tokens=256,
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True
-    )
+    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
+    output = model.generate(input_ids, max_new_tokens=256)
 
     reply = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)
     return reply
 
 # ---------------- VISION FUNCTION ----------------
-
 def run_vision(image, text_input, model_id):
     if model_id not in models:
         return "⚠️ Vision model not loaded."
@@ -103,34 +83,20 @@ def run_vision(image, text_input, model_id):
     template = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = processor(template, [img], return_tensors="pt")
 
-    output = model_vision.generate(
-        **inputs,
-        max_new_tokens=400,
-        do_sample=False,
-    )
-
+    output = model_vision.generate(**inputs, max_new_tokens=400, do_sample=False)
     output = output[:, inputs["input_ids"].shape[1]:]
-    text = processor.batch_decode(output, skip_special_tokens=True)[0]
-    return text
-
-# ---------------- FASTAPI BACKEND API ----------------
+    return processor.batch_decode(output, skip_special_tokens=True)[0]
 
+# ---------------- FASTAPI BACKEND ----------------
 api = FastAPI()
 
 @api.get("/health")
 def health():
-    return {
-        "status": "ok",
-        "device": device,
-        "chat_model": MODEL_ID,
-        "vision_loaded": len(models) > 0
-    }
+    return {"status": "ok", "device": device, "vision_loaded": len(models) > 0}
 
 @api.post("/api/chat")
 async def api_chat(message: str = Form(...)):
-    conversation = [{"role": "system", "content": "You are a helpful assistant."},
-                    {"role": "user", "content": message}]
-    
+    conversation = [{"role": "user", "content": message}]
     input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
     output = model.generate(input_ids, max_new_tokens=256)
     reply = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)
@@ -142,7 +108,6 @@ async def api_vision(image: UploadFile = File(...), text_input: str = Form("Desc
     return {"response": run_vision(np.array(img), text_input, model_id)}
 
 # ---------------- GRADIO UI ----------------
-
 def create_ui():
     with gr.Blocks() as demo:
         with gr.Tab("Chat"):
@@ -151,7 +116,7 @@ def create_ui():
         with gr.Tab("Vision"):
             img = gr.Image()
             txt = gr.Textbox("What's in this image?")
-            model_sel = gr.Dropdown(choices=list(models.keys()), value=list(models.keys())[0] if models else None)
+            model_sel = gr.Dropdown(choices=list(models.keys()))
             out = gr.Textbox()
             gr.Button("Analyze").click(run_vision, [img, txt, model_sel], out)