Upload 2 files

bot.py

@@ -7,6 +7,7 @@ Corre como una tarea en segundo plano dentro del mismo servidor FastAPI
 """
 import os
 import io
+import tempfile
 from PIL import Image
 import json
 import uuid
@@ -16,6 +17,7 @@ from datetime import date, datetime, timezone
 import httpx
 import firebase_admin
 from firebase_admin import credentials, firestore
+from gradio_client import Client, handle_file
 from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup, LabeledPrice
 from telegram.request import HTTPXRequest
 from telegram.ext import (
@@ -24,11 +26,12 @@ from telegram.ext import (
 )
 
 BOT_TOKEN = os.environ.get("TELEGRAM_BOT_TOKEN", "")
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
 
 REMOVE_BG_URL = "https://zgrafic-remove-bg-api.hf.space/remove-bg"
 ENHANCE_URL = "https://zgrafic-remove-bg-api.hf.space/enhance"
 VECTORIZE_URL = "https://zgrafic-remove-bg-api.hf.space/vectorize"
-VECTORIZE_AI_URL = "https://zgrafic-enhance-photo-api.hf.space/vectorize-ai"
+ZEROGPU_SPACE = "zgrafic/ia"  # Vectorizar con IA (HiVG-3B-Base + ZeroGPU)
 
 FREE_USES_PER_DAY = 10
 STARS_PER_IMAGE = 10
@@ -330,6 +333,45 @@ async def on_successful_payment(update: Update, context: ContextTypes.DEFAULT_TY
     await process_and_send(context, job["chat_id"], job["file_id"], job["mode"])
 
 
+_gradio_client = None
+
+
+def _get_gradio_client():
+    global _gradio_client
+    if _gradio_client is None:
+        _gradio_client = Client(ZEROGPU_SPACE, hf_token=HF_TOKEN or None)
+    return _gradio_client
+
+
+def _call_vectorize_ai_sync(image_bytes: bytes) -> str:
+    """Llama al Space de ZeroGPU (Gradio) — es una librería con llamadas
+    que bloquean, por eso se corre en un hilo aparte desde afuera."""
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+        tmp.write(image_bytes)
+        tmp_path = tmp.name
+
+    client = _get_gradio_client()
+    result_path = client.predict(handle_file(tmp_path), api_name="/predict")
+
+    with open(result_path, "r") as f:
+        return f.read()
+
+
+async def call_vectorize_ai(context: ContextTypes.DEFAULT_TYPE, chat_id: int, image_bytes: bytes):
+    """Vectoriza con HiVG-3B-Base, corriendo en el Space de ZeroGPU
+    (GPU real, no el servidor de CPU) — devuelve el SVG como texto, o
+    None si falló (ya le avisó el motivo a la persona en el chat)."""
+    try:
+        svg_text = await asyncio.to_thread(_call_vectorize_ai_sync, image_bytes)
+        return svg_text
+    except Exception as e:
+        await context.bot.send_message(
+            chat_id,
+            f"No se pudo vectorizar con IA ({type(e).__name__}): {e or 'sin detalle'}"
+        )
+        return None
+
+
 async def post_with_wakeup(context: ContextTypes.DEFAULT_TYPE, chat_id: int, url: str, files: dict, data: dict, mode: str):
     """Manda el pedido al servidor, reintentando con paciencia si está
     'despertando' (Hugging Face devuelve 429 mientras el contenedor
@@ -407,18 +449,23 @@ async def process_and_send(context: ContextTypes.DEFAULT_TYPE, chat_id: int, fil
             resp = await post_with_wakeup(context, chat_id, VECTORIZE_URL, files, {}, "vectorize")
             if resp is None:
                 return False
+            result_bytes = resp.content
+
+        elif mode == "vectorize_ai":
+            svg_text = await call_vectorize_ai(context, chat_id, image_bytes)
+            if svg_text is None:
+                return False
+            result_bytes = svg_text.encode("utf-8")
 
         else:
             if mode == "remove":
                 url = REMOVE_BG_URL
-            elif mode == "vectorize_ai":
-                url = VECTORIZE_AI_URL
             else:
                 url = ENHANCE_URL
 
             files = {"file": ("imagen.jpg", image_bytes, "image/jpeg")}
 
-            if mode in ("remove", "vectorize_ai"):
+            if mode == "remove":
                 data = {}
             elif mode == "enhance_imagen":
                 data = {"type": "imagen", "scale": "4"}
@@ -428,11 +475,12 @@ async def process_and_send(context: ContextTypes.DEFAULT_TYPE, chat_id: int, fil
             resp = await post_with_wakeup(context, chat_id, url, files, data, mode)
             if resp is None:
                 return False
+            result_bytes = resp.content
 
         extension = "svg" if mode in ("vectorize", "vectorize_ai") else "png"
         await context.bot.send_document(
             chat_id,
-            document=io.BytesIO(resp.content),
+            document=io.BytesIO(result_bytes),
             filename=f"resultado.{extension}",
         )
         return True

requirements.txt

@@ -18,3 +18,4 @@ realesrgan==0.3.0
 gfpgan==1.3.8
 firebase-admin==6.5.0
 vtracer==0.6.15
+gradio_client