Update app.py

app.py

@@ -3,52 +3,41 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
 from diffusers import StableDiffusionPipeline
 from PIL import Image
+import io
 
-# 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
+# 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.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
+    "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:
-        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
+        response = text_pipeline(input_text, max_new_tokens=150, do_sample=True)[0]['generated_text']
+        return response, None
 
-# Interfaz
+# Interfaz Gradio
 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])
+    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()