funciones.py

import io
import random
import globales
import herramientas
import gradio_client
import conexion_firebase
from datetime import datetime
from fastapi import HTTPException, status
from huggingface_hub import InferenceClient

servidor = globales.servidor
servicio = globales.servicio

def genera_platillo_gpu(platillo):

    prompt = globales.previo + platillo
    
    try:
        dict_espacios = conexion_firebase.obtenDato('nowme', servidor, 'espacios')

        espacios_habilitados = [
        nombre for nombre, config in dict_espacios.items()
        if config.get("habilitado", False) # Usamos .get() para evitar KeyError si 'habilitado' no existe
        ]

        print(f"Espacios habilitados: {espacios_habilitados}")
        
        espacio_aleatorio_elegido = random.choice(espacios_habilitados)
        configuracion_espacio = dict_espacios[espacio_aleatorio_elegido]
        client = gradio_client.Client(configuracion_espacio['ruta'], hf_token=globales.llave)
        
        result = client.predict(
        #**kwargs,
        prompt=prompt,
        #negative_prompt="live animals",
        seed=42,
        randomize_seed=False,
        width=786,
        height=568,
        # guidance_scale=3.5,
        # num_inference_steps=28,
        api_name=configuracion_espacio['api_name']       
        )

        #Cuando es GPU, debe de restar segundos disponibles de HF 
        herramientas.restaSegundosGPU(globales.work_cost)

        print("Platillo generado:", platillo)
        conexion_firebase.registraLog(servidor, platillo, espacio_aleatorio_elegido, 42)
        
        #Generales
        conexion_firebase.incrementar_campo_numerico('estadisticas', 'imagenes', 'imagenes_totales', amount=1)
        conexion_firebase.incrementar_campo_numerico('estadisticas', 'imagenes', 'endpoint-' + servicio, amount=1)
        conexion_firebase.incrementar_campo_numerico('estadisticas', 'imagenes', 'gpu-totales', amount=1)
        conexion_firebase.incrementar_campo_numerico('estadisticas', 'imagenes', 'gpu-' + servicio, amount=1)
        #Día a Día
        now = datetime.now() 
        document_id_fecha_hora = now.strftime("%Y-%m-%d")
        nombre_doc = document_id_fecha_hora + " - " +  servidor 
        conexion_firebase.incrementar_campo_numerico('estadisticas', nombre_doc, 'imagenes_totales', amount=1)
        conexion_firebase.incrementar_campo_numerico('estadisticas', nombre_doc, 'endpoint-' + servicio, amount=1)
        conexion_firebase.incrementar_campo_numerico('estadisticas', nombre_doc, 'gpu-totales', amount=1)
        conexion_firebase.incrementar_campo_numerico('estadisticas', nombre_doc, 'gpu-' + servicio, amount=1)    

        return result[0]

    except Exception as e:     
        print("Excepción: ", e)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=e
        )

def genera_platillo_inference(platillo):
    
    dict_modelos = conexion_firebase.obtenDato('nowme', servidor, 'modelos')    
    modelos_habilitados = [
        nombre for nombre, config in dict_modelos.items()
        if config.get("habilitado", False) # Usamos .get() para evitar KeyError si 'habilitado' no existe
        ]
    
    print(f"Modelos habilitados: {modelos_habilitados}")
    modelo_aleatorio_elegido = random.choice(modelos_habilitados)
    configuracion_modelo = dict_modelos[modelo_aleatorio_elegido]   

    creditos_restantes_inference = conexion_firebase.obtenDato('nowme', servidor, 'inferencias')
    
    if creditos_restantes_inference > 0:
        provedor_seleccionado = globales.proveedor
    else:
        provedor_seleccionado = globales.proveedor_back

    prompt = globales.previo + platillo
    print("Platillo enviado:", platillo)
       
    client = InferenceClient(
            provider= provedor_seleccionado,
            model=configuracion_modelo['ruta'],
            api_key=globales.llave
        )

    try: 
        image = client.text_to_image(
        prompt,
        #negative_prompt="live animals",        
        width=784, #786
        height=560, #568
        num_inference_steps=12,
        seed=42
        )

        herramientas.restaSegundosInference(globales.inference_cost)              
        
    except Exception as e:
        print("Excepción: ", e)
        if "Gateway Time-out" in str(e): 
            print("GATEWAY TIME-OUT 💀")
                   
        error_impreso = f"Error: {e}"
        print(error_impreso)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        )
    
    img_io = io.BytesIO()
    image.save(img_io, "PNG")
    img_io.seek(0)
    print("Platillo generado:", platillo)
    conexion_firebase.registraLog(servidor, platillo, modelo_aleatorio_elegido, 42)
    #Generales
    conexion_firebase.incrementar_campo_numerico('estadisticas', 'imagenes', 'imagenes_totales', amount=1)
    conexion_firebase.incrementar_campo_numerico('estadisticas', 'imagenes', 'endpoint-' + servicio, amount=1)
    conexion_firebase.incrementar_campo_numerico('estadisticas', 'imagenes', 'inference-totales', amount=1)
    conexion_firebase.incrementar_campo_numerico('estadisticas', 'imagenes', 'inference-' + servicio, amount=1)
    #Día a Día
    now = datetime.now() 
    document_id_fecha_hora = now.strftime("%Y-%m-%d")
    nombre_doc = document_id_fecha_hora + " - " +  servidor   
    conexion_firebase.incrementar_campo_numerico('estadisticas', nombre_doc, 'imagenes_totales', amount=1)
    conexion_firebase.incrementar_campo_numerico('estadisticas', nombre_doc, 'endpoint-' + servicio, amount=1)
    conexion_firebase.incrementar_campo_numerico('estadisticas', nombre_doc, 'inference-totales', amount=1)
    conexion_firebase.incrementar_campo_numerico('estadisticas', nombre_doc, 'inference-' + servicio, amount=1)

    return img_io