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

# Modelo de texto
text_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(text_model_name)
text_model = AutoModelForCausalLM.from_pretrained(
    text_model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)
text_pipe = TextGenerationPipeline(
    model=text_model,
    tokenizer=tokenizer,
    max_new_tokens=200,
    do_sample=True,
    temperature=0.8,
    top_p=0.95
)

# Modelo de imagen
image_pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
).to("cuda")

# Lógica para decidir si es imagen o texto
def is_image_prompt(prompt):
    keywords = ["dibuja", "genera una imagen", "imagen de", "píntame", "crea una ilustración"]
    return any(kw in prompt.lower() for kw in keywords)

# Función del bot
def bot_response(message):
    if is_image_prompt(message):
        image = image_pipe(message).images[0]
        return "", image
    else:
        prompt = "Eres una asistente coqueta, creativa y dulce.\nUsuario: " + message + "\nAsistente:"
        result = text_pipe(prompt)[0]['generated_text']
        reply = result.split("Asistente:")[-1].strip()
        return reply, None

# Interfaz
with gr.Blocks() as demo:
    gr.Markdown("## Asistente inteligente de texto e imágenes")
    input_box = gr.Textbox(label="Tu mensaje", placeholder="Escribe lo que quieras...")
    text_output = gr.Textbox(label="Respuesta de texto")
    image_output = gr.Image(label="Imagen generada")
    input_box.submit(fn=bot_response, inputs=input_box, outputs=[text_output, image_output])

demo.launch()