app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os

# Obtén el token de manera segura desde el entorno
hf_token = os.getenv("HF_API_TOKEN")

# Clase para manejar múltiples modelos
class ModelHandler:
    def __init__(self, model_names, token):
        """
        Inicializa el manejador de modelos con los nombres de los modelos y el token de API.
        """
        self.clients = {
            model_name: InferenceClient(model_name, token=token)
            for model_name in model_names
        }
        self.current_model = model_names[0]

    def switch_model(self, model_name):
        """
        Cambia el modelo actual.
        """
        if model_name in self.clients:
            self.current_model = model_name
        else:
            raise ValueError(f"Modelo {model_name} no está disponible.")

    def generate_response(self, input_text):
        """
        Genera una respuesta utilizando el modelo actual.
        """
        prompt = f"Debes de responder a cualquier pregunta:\nPregunta: {input_text}"
        try:
            messages = [{"role": "user", "content": prompt}]
            client = self.clients[self.current_model]
            response = client.chat_completion(messages=messages, max_tokens=500)
            if hasattr(response, 'choices') and response.choices:
                return response.choices[0].message.content
            else:
                return str(response)
        except Exception as e:
            return f"Error al realizar la inferencia: {e}"

# Lista de modelos disponibles
model_names = [
    "microsoft/Phi-3-mini-4k-instruct"
]

# Inicializa el manejador de modelos
model_handler = ModelHandler(model_names, hf_token)

# Define la función para generación de imágenes
def generate_image(prompt):
    """
    Genera una imagen utilizando un modelo de generación de imágenes.
    """
    try:
        client = InferenceClient("CompVis/stable-diffusion-v1-4", token=hf_token)
        image = client.text_to_image(prompt, width=512, height=512)
        return image
    except Exception as e:
        return f"Error al generar la imagen: {e}"

# Configura la interfaz en Gradio con selección de modelos y generación de imágenes
with gr.Blocks(title="Multi-Model LLM Chatbot with Image Generation") as demo:
    gr.Markdown(
        """
        ## Chatbot Multi-Modelo LLM con Generación de Imágenes
        Este chatbot permite elegir entre múltiples modelos de lenguaje para responder preguntas o generar imágenes 
        a partir de descripciones.
        """
    )
    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=model_names + ["Generación de Imágenes"],
            value=model_names[0],
            label="Seleccionar Acción/Modelo",
            interactive=True
        )
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                lines=5,
                placeholder="Escribe tu consulta o descripción para la imagen...",
                label="Entrada"
            )
        with gr.Column():
            output_display = gr.Textbox(
                lines=5,
                label="Respuesta",
                interactive=False,
                visible=False
            )
            output_image = gr.Image(
                label="Imagen Generada",
                interactive=False,
                visible=False
            )
    submit_button = gr.Button("Enviar")
    
    # Define la función de actualización
    def process_input(selected_action, user_input):
        if selected_action == "Generación de Imágenes":
            return "", generate_image(user_input), gr.update(visible=True), gr.update(visible=False)
        else:
            model_handler.switch_model(selected_action)
            response = model_handler.generate_response(user_input)
            return response, None, gr.update(visible=False), gr.update(visible=True)
    
    # Conecta la función a los componentes
    submit_button.click(
        fn=process_input,
        inputs=[model_dropdown, input_text],
        outputs=[output_display, output_image, output_image, output_display]
    )

# Lanza la interfaz
demo.launch()