Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException # <--- AÑADIDO HTTPException
+from fastapi import FastAPI, HTTPException
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
 from fastapi.responses import Response
@@ -19,7 +19,7 @@ import google.generativeai as genai
 from typing import Optional, List, Any, Dict, Union
 from diffusers import StableDiffusionPipeline, LCMScheduler
 
-app = FastAPI(title="Mirage Medical Search API - Lite Version")
+app = FastAPI(title="MIRAGE")
 
 app.add_middleware(
     CORSMiddleware,
@@ -29,22 +29,22 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- 1. CONFIGURACIÓN DE MODELOS ---
+# Models
 MODEL_NAME = 'hf-hub:luhuitong/CLIP-ViT-L-14-448px-MedICaT-ROCO'
 HF_DATASET_ID = "mdwiratathya/ROCO-radiology"
 SPLIT = "train"
 device = "cpu"
 
-# Variables Globales
+# Glob variables
 model = None
 tokenizer = None
-embeddings = None        # Image Embeddings (Visual Only)
+embeddings = None       
 metadata = None
 dataset_stream = None 
 gemini_available = False
 pipe = None 
 
-# --- 2. AUTENTICACIÓN ---
+# authentication
 try:
     hf_token = os.environ.get('HF_TOKEN')
     if hf_token:
@@ -58,7 +58,7 @@ try:
 except Exception as e:
     print(f"Error auth: {e}")
 
-# --- HELPER: PLACEHOLDER ---
+# to handle if there's an error
 def create_placeholder_image(text="Image Error"):
     img = Image.new('RGB', (512, 512), color=(40, 40, 45))
     d = ImageDraw.Draw(img)
@@ -73,24 +73,15 @@ def create_placeholder_image(text="Image Error"):
     img.save(img_byte_arr, format='JPEG')
     return img_byte_arr.getvalue()
 
-# --- 3. CARGA DE DATOS ---
+# load the data
 @app.on_event("startup")
 async def load_data():
     global model, tokenizer, embeddings, metadata, dataset_stream, pipe
-    print("--- INICIANDO MIRAGE BACKEND (Lite Version) ---")
-    
-    # 1. CARGAR CLIP
-    try:
-        print("👁️ Cargando CLIP...")
-        model, _, _ = open_clip.create_model_and_transforms(MODEL_NAME, device=device)
-        tokenizer = open_clip.get_tokenizer(MODEL_NAME)
-        model.eval()
-        print("✅ CLIP cargado.")
-    except Exception as e:
-        print(f"❌ Error CLIP: {e}")
+    model, _, _ = open_clip.create_model_and_transforms(MODEL_NAME, device=device)
+    tokenizer = open_clip.get_tokenizer(MODEL_NAME)
+    model.eval()
 
-    # 2. CARGAR METADATA
-    print("📦 Cargando Metadata...")
+    # load metadata
     if os.path.exists("metadata_text.json"):
         with open("metadata_text.json", 'r') as f:
             metadata = json.load(f)
@@ -98,53 +89,34 @@ async def load_data():
         with open("metadata.json", 'r') as f:
             metadata = json.load(f)
     else:
-        print("⚠️ NO SE ENCONTRÓ METADATA.")
+        print("no metadata file found")
         metadata = [{"dataset_index": 0, "filename": "error", "caption": "Error"}]
 
-    # 3. CARGAR EMBEDDINGS DE IMAGEN (SOLO IMAGEN)
-    if os.path.exists("embeddings.npy"):
-        embeddings = np.load("embeddings.npy")
-        print(f"✅ Image Embeddings listos: {embeddings.shape[0]} registros.")
+    # load the embdeddings of the images (already processed)
+    embeddings = np.load("embeddings.npy")
+    print(f"✅ Image Embeddings listos: {embeddings.shape[0]} registros.")
+
+    # load the dataset
+    dataset_stream = load_dataset(HF_DATASET_ID, split=SPLIT, streaming=False) 
+
+    # Load the Stable Diffusion LCM Pipeline
+    model_id = "Nihirc/Prompt2MedImage"
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+    pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
+    pipe.fuse_lora() 
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, solver_order=2)
+    pipe.safety_checker = None
+    pipe.requires_safety_checker = False
+    
+    if device == "cpu":
+        pipe = pipe.to("cpu")
+        pipe.enable_attention_slicing() 
     else:
-        print("⚠️ NO SE ENCONTRARON IMAGE EMBEDDINGS.")
-        embeddings = np.zeros((1, 768))
-
-    # 4. CARGAR DATASET
-    try:
-        print("📦 Cargando Dataset en RAM (1-2 mins)...")
-        dataset_stream = load_dataset(HF_DATASET_ID, split=SPLIT, streaming=False) 
-        print(f"✅ Dataset listo. Total: {len(dataset_stream)}")
-    except Exception as e:
-        print(f"❌ Error dataset: {e}")
-        dataset_stream = None
-
-    # 5. CARGAR STABLE DIFFUSION (LCM)
-    print("🎨 Cargando modelo generativo (LCM Mode)...")
-    try:
-        model_id = "Nihirc/Prompt2MedImage"
-        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
-        print("⚡ Inyectando pesos LCM-LoRA...")
-        pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
-        pipe.fuse_lora() 
-        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, solver_order=2)
-        pipe.safety_checker = None
-        pipe.requires_safety_checker = False
-        
-        if device == "cpu":
-            pipe = pipe.to("cpu")
-            pipe.enable_attention_slicing() 
-        else:
-            pipe = pipe.to("cuda")
-        print("✅ Generador LCM listo.")
-    except Exception as e:
-        print(f"❌ Error Generador: {e}")
-
-
-# --- 4. FUNCIONES CORE ---
+        pipe = pipe.to("cuda")
 
 def calculate_vector(text, add=None, sub=None):
     with torch.no_grad():
-        # Usamos el texto tal cual viene del usuario
+        # the user gives us a text, we obtain the embedding using CLIP
         text_tokens = tokenizer([text]).to(device)
         vec = model.encode_text(text_tokens)
         vec /= vec.norm(dim=-1, keepdim=True)
@@ -160,17 +132,12 @@ def calculate_vector(text, add=None, sub=None):
         return vec
 
 def get_retrieval_and_context(query_vector, top_k):
-    """
-    Realiza el retrieval basado EXCLUSIVAMENTE en similitud visual.
-    Query Text Vector vs Image Embeddings.
-    """
+    # We compare the query (text) embd with the image embeddings to retrieve
     query_vec_np = query_vector.cpu().numpy()
     
-    # Similitud Visual (Query vs Image Embeddings)
-    # query_vec_np es (1, 768), embeddings es (N, 768) -> resultado (N,)
+   
+    # query_vec_np (1, 768), embeddings (N, 768) -> result (N,)
     sim_img = (query_vec_np @ embeddings.T).squeeze()
-    
-    # Ordenar índices (descendente)
     best_indices = sim_img.argsort()[-top_k:][::-1]
     
     real_matches = []
@@ -185,7 +152,7 @@ def get_retrieval_and_context(query_vector, top_k):
         
         real_matches.append({
             "url": f"/image/{safe_index}",
-            "score": float(sim_img[idx]), # Solo score visual
+            "score": float(sim_img[idx]), 
             "filename": meta.get("filename", "img"),
             "caption": meta.get("caption", ""),
             "index": safe_index 
@@ -202,7 +169,6 @@ def generate_llm_prompt(captions, user_text):
         return user_text + ". " + (captions[0] if captions else "")
     try:
         llm = genai.GenerativeModel('gemini-2.5-flash')
-        # Prompt actualizado para usar directamente el texto del usuario
         prompt = f"Using the following medical query: '{user_text}', synthesize these findings into a concise radiology description: {', '.join(captions[:3])}"
         res = llm.generate_content(prompt)
         return res.text.strip()
@@ -217,7 +183,7 @@ def generate_synthetic_image(prompt, steps=5, guidance=1.5):
         image = pipe(prompt[:77], height=512, width=512, num_inference_steps=steps, guidance_scale=guidance, negative_prompt=NEGATIVE_PROMPT).images[0]
         
         draw = ImageDraw.Draw(image)
-        text = "Created by MIRAGE LITE"
+        text = "Created by MIRAGE"
         try: font = ImageFont.load_default() 
         except: font = None
         bbox = draw.textbbox((0, 0), text, font=font)
@@ -239,8 +205,8 @@ def fetch_image_from_stream(index):
         return dataset_stream[idx]['image']
     except Exception: return None
 
-# --- ENDPOINTS ---
-# CAMBIO IMPORTANTE: Renombrado a /api/health para dejar libre la raíz "/"
+
+# ENDPOINTS
 @app.get("/api/health")
 def health_check(): 
     return {"status": "online", "version": "lite"}
@@ -261,7 +227,6 @@ def get_image(index: str):
     except Exception: pass
     return Response(content=create_placeholder_image("Error"), media_type="image/jpeg")
 
-# --- MODELOS PYDANTIC SIMPLIFICADOS ---
 class GenerationRequest(BaseModel):
     original_text: str
     sub_concept: Optional[str] = None
@@ -272,17 +237,14 @@ class GenerationRequest(BaseModel):
     guidance_scale: float = 1.5   
     num_inference_steps: int = 5
 
-# --- ENDPOINT PRINCIPAL ---
+# this is the main endpoint
 @app.post("/generate_comparison")
 def generate_comparison(req: GenerationRequest):
-    if not model: raise HTTPException(status_code=503, detail="Loading...") # HTTPException ahora funciona
+    if not model: raise HTTPException(status_code=503, detail="Loading...") 
     try:
-        # ASIGNACIÓN DIRECTA SIN TRADUCCIÓN
         final_query = req.original_text
         final_add = req.add_concept
         final_sub = req.sub_concept
-            
-        print(f"⚡ Procesando Lite (Raw Input): '{final_query}'")
         
         response_data = {
             "original_text": final_query,
@@ -292,13 +254,11 @@ def generate_comparison(req: GenerationRequest):
             "input_lang_detected": "raw"
         }
 
-        # 1. PROCESAR ORIGINAL (Siempre Visual Search)
         vec_orig = calculate_vector(final_query)
         match_orig, caps_orig = get_retrieval_and_context(vec_orig, req.top_k)
         
         prompt_orig = ""
         if req.gen_text:
-            # Pasa el texto original al LLM
             prompt_orig = generate_llm_prompt(caps_orig, final_query)
         else:
             prompt_orig = "LLM generation skipped."
@@ -316,7 +276,6 @@ def generate_comparison(req: GenerationRequest):
             }
         }
 
-        # 2. PROCESAR MODIFICADO (Dual Search - Aritmética)
         has_dual = (final_add and final_add.strip()) and (final_sub and final_sub.strip())
         if has_dual:
             vec_mod = calculate_vector(final_query, final_add, final_sub)
@@ -324,7 +283,6 @@ def generate_comparison(req: GenerationRequest):
             
             prompt_mod = ""
             if req.gen_text:
-                # Construye el string de aritmética sin traducción
                 prompt_mod = generate_llm_prompt(caps_mod, f"{final_query} + {final_add} - {final_sub}")
             else:
                 prompt_mod = "LLM generation skipped."
@@ -354,27 +312,20 @@ def search(req: GenerationRequest):
     return generate_comparison(req)
 
 
-# --- SERVIR FRONTEND ---
-
-# 1. Montar los assets estáticos (JS, CSS que genera Vite)
+# To create the frontend serving
 app.mount("/assets", StaticFiles(directory="static/assets"), name="assets")
 
-# 2. Servir imágenes si las hay en public
 if os.path.exists("static/images"):
     app.mount("/images", StaticFiles(directory="static/images"), name="images")
 
-# 3. Ruta raíz -> Devuelve el HTML principal
 @app.get("/")
 async def read_index():
     return FileResponse('static/index.html')
 
-# 4. Catch-all: Cualquier otra ruta devuelve index.html (para que React Router no falle al recargar)
 @app.get("/{full_path:path}")
 async def catch_all(full_path: str):
-    # Si intentan pedir un archivo que existe (ej. un .png), lo damos
     if os.path.exists(f"static/{full_path}"):
         return FileResponse(f"static/{full_path}")
-    # Si no, devolvemos la app de React
     return FileResponse('static/index.html')
 
 if __name__ == "__main__":