import streamlit as st
from diffusers import DiffusionPipeline
import requests

# Simulación de las APIs (sustituir por las URLs reales de las APIs)
LLAMA_API_URL = "https://api-inference.huggingface.co/models/meta-llama/Llama-3.2-11B-Vision-Instruct"
FLUX_API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
COGVIDEO_API_URL = DiffusionPipeline.from_pretrained("bertjiazheng/KoolCogVideoX-5b")


st.title("ChatBot Multi-IA")

# Sección para chat y análisis de imágenes (meta-llama/Llama-3.2-11B-Vision-Instruct)
st.header("Chat y análisis de imágenes")

# Entrada de texto para chat
user_input = st.text_area("Escribe tu pregunta o sube una imagen para analizar", "")

# Entrada de imagen para análisis
uploaded_image = st.file_uploader("Sube una imagen para que la IA la analice", type=["png", "jpg", "jpeg"])

# Botón para enviar la consulta
if st.button("Enviar para análisis"):
    if uploaded_image:
        # Enviar la imagen a la API de análisis de imágenes
        files = {"image": uploaded_image.getvalue()}
        response = requests.post(LLAMA_API_URL, files=files)
        st.image(uploaded_image, caption="Imagen subida", use_column_width=True)
        st.write("Respuesta de análisis:", response.text)
    elif user_input:
        # Enviar el texto al modelo Llama para chat
        data = {"text": user_input}
        response = requests.post(LLAMA_API_URL, json=data)
        st.write("Respuesta del chat:", response.text)
    else:
        st.warning("Escribe algo o sube una imagen para análisis")

# Sección para generación de imágenes (black-forest-labs/FLUX.1-dev)
st.header("Generación de imágenes")

# Entrada de texto para la generación de imagen
image_prompt = st.text_input("Describe la imagen que deseas generar", "")

# Botón para generar la imagen
if st.button("Generar imagen"):
    if image_prompt:
        # Enviar la descripción al modelo FLUX.1-dev
        data = {"prompt": image_prompt}
        response = requests.post(FLUX_API_URL, json=data)
        st.image(response.content, caption="Imagen generada", use_column_width=True)
    else:
        st.warning("Por favor, ingresa una descripción para generar una imagen.")

# Sección para generación de video (THUDM/CogVideoX-5b)
st.header("Generación de video")

# Entrada de texto para la generación de video
video_prompt = st.text_input("Describe el video que deseas generar", "")

# Botón para generar el video
if st.button("Generar video"):
    if video_prompt:
        # Enviar la descripción al modelo CogVideoX-5b
        data = {"prompt": video_prompt}
        response = requests.post(COGVIDEO_API_URL, json=data)
        # Guardar y mostrar el video generado
        with open("video_generado.mp4", "wb") as f:
            f.write(response.content)
        st.video("video_generado.mp4")
    else:
        st.warning("Por favor, ingresa una descripción para generar un video.")

# Footer
st.write("Aplicación que combina chat, generación de imágenes y videos con múltiples IAs.")