128 lines
5.0 KiB
Python
128 lines
5.0 KiB
Python
import os
|
||
import requests
|
||
import logging
|
||
import json
|
||
|
||
# Configuración de logging
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
||
|
||
# Configuración de APIs
|
||
HF_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
|
||
HF_MODEL = "stabilityai/stable-diffusion-3.5-large"
|
||
#HF_MODEL = "stabilityai/stable-diffusion-2-1" # Modelo de Stable Diffusion
|
||
OLLAMA_API_URL = "http://10.233.23.140:11434/api/generate" # URL de tu instancia de Ollama
|
||
|
||
# Configuración del servidor Bepasty
|
||
PASTE_SERVER = "http://10.233.37.160"
|
||
PASTE_USERNAME = os.getenv("PASTE_USERNAME")
|
||
PASTE_PASSWORD = os.getenv("PASTE_PASSWORD") # Debería almacenarse en un Secret de Kubernetes
|
||
|
||
# Ruta temporal para la imagen generada
|
||
TEMP_IMAGE_PATH = "/tmp/generated_image"
|
||
|
||
def translate_to_english(text):
|
||
"""Traduce el texto al inglés usando la instancia de Ollama y maneja respuestas en streaming."""
|
||
payload = {
|
||
"model": "gemma2:2b",
|
||
"prompt": f"Traduce al inglés UNA sola versión, lo mas sugerente posible sin agreagr titulos o texto superfluo solo la traduccion: {text}",
|
||
"stream": False # Aseguramos que Ollama no use streaming
|
||
}
|
||
|
||
try:
|
||
response = requests.post(OLLAMA_API_URL, json=payload, timeout=120)
|
||
response.raise_for_status()
|
||
|
||
response_json = response.json()
|
||
|
||
# Si Ollama responde correctamente, extraer el texto
|
||
return response_json.get("response", text).strip()
|
||
|
||
except requests.exceptions.Timeout:
|
||
logging.error("⏳ Ollama tardó demasiado en responder. Usando texto original.")
|
||
return text
|
||
except requests.exceptions.RequestException as e:
|
||
logging.error(f" Error en la traducción con Ollama: {e}")
|
||
return text
|
||
|
||
def get_paste_token():
|
||
"""Obtiene un token de autenticación en paste.priet.us"""
|
||
url = f"{PASTE_SERVER}/api/token"
|
||
payload = {"username": PASTE_USERNAME, "password": PASTE_PASSWORD}
|
||
|
||
try:
|
||
response = requests.post(url, json=payload, timeout=10)
|
||
response.raise_for_status()
|
||
return response.json().get("token", "Error: No se pudo autenticar en paste.priet.us.")
|
||
except requests.exceptions.RequestException as e:
|
||
logging.error(f"Error al obtener token: {e}")
|
||
return None
|
||
|
||
def upload_to_paste(image_path):
|
||
"""Sube la imagen generada y devuelve la URL pública"""
|
||
token = get_paste_token()
|
||
if not token:
|
||
return "Error: No se pudo obtener el token de autenticación."
|
||
|
||
url = f"{PASTE_SERVER}/paste"
|
||
headers = {"Authorization": f"Bearer {token}"}
|
||
|
||
with open(image_path, "rb") as image_file:
|
||
files = {"c": image_file}
|
||
try:
|
||
response = requests.post(url, headers=headers, files=files, data={"expire": "yes", "private": "no"})
|
||
response.raise_for_status()
|
||
paste_data = response.json()
|
||
paste_url = paste_data.get("url", "Error: No se pudo obtener la URL.")
|
||
return paste_url.replace("http://10.233.37.160", "https://paste.priet.us")
|
||
except requests.exceptions.RequestException as e:
|
||
return f"Error al subir la imagen: {e}"
|
||
|
||
def generate_image(prompt):
|
||
"""Genera una imagen con la API de Hugging Face y la sube a paste.priet.us"""
|
||
if not HF_API_KEY:
|
||
return "Error: No se encontró la API Key de Hugging Face."
|
||
|
||
translated_prompt = translate_to_english(prompt)
|
||
|
||
url = f"https://api-inference.huggingface.co/models/{HF_MODEL}"
|
||
headers = {"Authorization": f"Bearer {HF_API_KEY}"}
|
||
payload = {"inputs": translated_prompt}
|
||
|
||
try:
|
||
logging.info(f"️ Generando imagen con prompt traducido: {translated_prompt}")
|
||
response = requests.post(url, headers=headers, json=payload, timeout=120)
|
||
|
||
if response.status_code != 200:
|
||
logging.error(f" Error en la generación de la imagen: {response.text}")
|
||
return f"Error en la generación de la imagen: {response.text}"
|
||
|
||
content_type = response.headers.get("Content-Type", "")
|
||
image_ext = "png" if "png" in content_type else "jpg" if "jpeg" in content_type else None
|
||
|
||
if not image_ext:
|
||
return "Error: No se recibió una imagen válida."
|
||
|
||
image_path = f"{TEMP_IMAGE_PATH}.{image_ext}"
|
||
|
||
with open(image_path, "wb") as f:
|
||
f.write(response.content)
|
||
|
||
paste_url = upload_to_paste(image_path)
|
||
os.remove(image_path) # Eliminamos la imagen después de subirla
|
||
return paste_url
|
||
except requests.exceptions.RequestException as e:
|
||
logging.error(f" Error al generar la imagen: {e}")
|
||
return f"Error al generar la imagen: {e}"
|
||
|
||
def run(sender, *args):
|
||
"""Ejecuta el comando de generación de imágenes."""
|
||
if not args:
|
||
return "Uso: .image <descripción>"
|
||
|
||
prompt = " ".join(args).strip()
|
||
return generate_image(prompt)
|
||
|
||
def help():
|
||
return "Uso: .image <descripción> - Genera una imagen con Stable Diffusion y la sube a paste.priet.us."
|
||
|