Update app.py

app.py

@@ -1,16 +1,19 @@
 import io
 import torch
+import numpy as np
 from fastapi import FastAPI, UploadFile, File, Form
 from fastapi.responses import StreamingResponse
 from PIL import Image
+
+# Importaciones de IA
 from diffusers import StableDiffusionPipeline, DDIMScheduler
-from diffusers.utils import load_image
+from huggingface_hub import hf_hub_download
 from ip_adapter.ip_adapter_faceid import IPAdapterFaceID
 from insightface.app import FaceAnalysis
 
 app = FastAPI(title="IP-Adapter-FaceID API")
 
-# Variables globales para los modelos
+# Variables globales para evitar recargar los modelos en cada petición
 app_insight = None
 ip_model = None
 
@@ -19,13 +22,16 @@ def load_models():
     global app_insight, ip_model
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Iniciando servidor. Hardware detectado: {device.upper()}")
     
     # 1. Inicializar InsightFace para detectar rostros
+    print("Cargando InsightFace...")
     app_insight = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
     app_insight.prepare(ctx_id=0, det_size=(640, 640))
     
     # 2. Inicializar Stable Diffusion v1.5
-    base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE" # O el modelo base que prefieras
+    print("Cargando modelo base de Stable Diffusion...")
+    base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
     noise_scheduler = DDIMScheduler(
         num_train_timesteps=1000,
         beta_start=0.00085,
@@ -42,13 +48,21 @@ def load_models():
         scheduler=noise_scheduler
     ).to(device)
     
-    # 3. Cargar las capas de IP-Adapter FaceID
-    ip_ckpt = "h94/IP-Adapter-FaceID/ip-adapter-faceid_sd15.bin"
-    ip_model = IPAdapterFaceID(pipe, ip_ckpt, device, num_tokens=4)
+    # 3. Descargar y cargar las capas de IP-Adapter FaceID
+    print("Descargando pesos de IP-Adapter-FaceID...")
+    ip_ckpt_path = hf_hub_download(
+        repo_id="h94/IP-Adapter-FaceID", 
+        filename="ip-adapter-faceid_sd15.bin", 
+        repo_type="model"
+    )
+    
+    print("Acoplando IP-Adapter al modelo base...")
+    ip_model = IPAdapterFaceID(pipe, ip_ckpt_path, device, num_tokens=4)
+    print("¡Todos los modelos cargados exitosamente!")
 
 @app.get("/")
 def read_root():
-    return {"status": "running", "info": "Envía un POST a /predict para generar imágenes"}
+    return {"status": "running", "message": "API de IP-Adapter-FaceID lista. Realiza un POST a /predict"}
 
 @app.post("/predict")
 async def predict(
@@ -57,34 +71,37 @@ async def predict(
     num_inference_steps: int = Form(30),
     face_image: UploadFile = File(...)
 ):
-    # Leer la imagen enviada desde el cliente (Kotlin, etc.)
-    face_bytes = await face_image.read()
-    image = Image.open(io.BytesIO(face_bytes)).convert("RGB")
-    
-    # Extraer el rostro (Embedding) usando InsightFace
-    import numpy as np
-    cv_img = np.array(image)
-    faces = app_insight.get(cv_img)
-    
-    if len(faces) == 0:
-        return {"error": "No se detectó ningún rostro en la imagen provista."}
+    try:
+        # 1. Leer y preparar la imagen
+        face_bytes = await face_image.read()
+        image = Image.open(io.BytesIO(face_bytes)).convert("RGB")
+        cv_img = np.array(image)
         
-    face_id_embed = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0).unsqueeze(0)
-    
-    # Generar la nueva imagen mezclando el prompt + identidad del rostro
-    images = ip_model.generate(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        face_id_embed=face_id_embed,
-        num_samples=1,
-        num_inference_steps=num_inference_steps,
-        seed=42
-    )
-    
-    # Convertir la imagen generada a bytes para retornarla por HTTP
-    output_image = images[0]
-    img_byte_arr = io.BytesIO()
-    output_image.save(img_byte_arr, format='PNG')
-    img_byte_arr.seek(0)
-    
-    return StreamingResponse(img_byte_arr, media_type="image/png")
+        # 2. Extraer el rostro (Embedding) usando InsightFace
+        faces = app_insight.get(cv_img)
+        
+        if len(faces) == 0:
+            return {"error": "No se detectó ningún rostro en la imagen provista. Por favor intenta con otra foto."}
+            
+        face_id_embed = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0).unsqueeze(0)
+        
+        # 3. Generar la imagen
+        images = ip_model.generate(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            face_id_embed=face_id_embed,
+            num_samples=1,
+            num_inference_steps=num_inference_steps,
+            seed=42 # Cambia o recibe este parámetro si quieres resultados aleatorios
+        )
+        
+        # 4. Formatear la respuesta como PNG
+        output_image = images[0]
+        img_byte_arr = io.BytesIO()
+        output_image.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+        
+        return StreamingResponse(img_byte_arr, media_type="image/png")
+        
+    except Exception as e:
+        return {"error": str(e)}