import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
from diffusers import StableDiffusionPipeline
from PIL import Image
import io

# Modelo de texto en CPU
text_model = "tiiuae/falcon-rw-1b"
tokenizer = AutoTokenizer.from_pretrained(text_model)
model = AutoModelForCausalLM.from_pretrained(text_model)
text_pipeline = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=-1)

# Modelo de imagen en CPU
image_pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", 
    torch_dtype=torch.float32
).to("cpu")

# Lógica para detectar si el prompt es de texto o imagen
def chatbot(input_text):
    if any(word in input_text.lower() for word in ["imagen", "dibuja", "pinta", "foto", "muestra"]):
        image = image_pipe(input_text).images[0]
        return None, image
    else:
        response = text_pipeline(input_text, max_new_tokens=150, do_sample=True)[0]['generated_text']
        return response, None

# Interfaz Gradio
with gr.Blocks() as demo:
    gr.Markdown("## Bot Generador de Texto e Imágenes (CPU)")

    with gr.Row():
        textbox = gr.Textbox(placeholder="Escribe algo... (ej: Dibuja una chica en la playa)")
        send = gr.Button("Enviar")

    with gr.Row():
        text_output = gr.Textbox(label="Respuesta de texto")
        image_output = gr.Image(label="Imagen generada")

    send.click(fn=chatbot, inputs=textbox, outputs=[text_output, image_output])

demo.launch()