Update app.py

app.py

@@ -1,36 +1,54 @@
-import gradio as gr import torch from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline from diffusers import StableDiffusionPipeline from PIL import Image import io import os
-
-Carga del modelo de texto
-
-text_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct") text_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instruct", device_map="auto", torch_dtype=torch.bfloat16) text_pipe = pipeline("text-generation", model=text_model, tokenizer=text_tokenizer)
-
-Carga del modelo de imágenes
-
-image_pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
-
-Memoria por usuario
-
-user_memory = {}
-
-def is_prompt_image(text): image_keywords = ["dibuja", "imagen", "pinta", "crea una imagen", "genera una foto", "haz un dibujo", "quiero ver"] return any(kw in text.lower() for kw in image_keywords)
-
-def build_prompt(user_id, message): personality = ( "Eres una asistente virtual sensual, dulce y coqueta. " "Respondes de forma cercana, creativa y provocadora. Usa emoticonos como 😘😉🔥.\n\n" ) if user_id not in user_memory: user_memory[user_id] = []
-
-user_memory[user_id].append(f"Usuario: {message}")
-user_memory[user_id] = user_memory[user_id][-5:]
-
-context = "\n".join(user_memory[user_id])
-return personality + context + "\nAsistente:"
-
-def responder(input_text, user_id="usuario1"): if is_prompt_image(input_text): image = image_pipe(input_text).images[0] return None, image else: prompt = build_prompt(user_id, input_text) result = text_pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.8, top_p=0.95)[0]['generated_text'] reply = result.split("Asistente:")[-1].strip() user_memory[user_id].append(f"Asistente: {reply}") return reply, None
-
-with gr.Blocks() as demo: with gr.Row(): chatbot = gr.Textbox(label="Escribe algo...") btn = gr.Button("Enviar") with gr.Row(): output_text = gr.Textbox(label="Respuesta de texto") output_image = gr.Image(label="Imagen generada")
-
-def on_click(user_input):
-    text, image = responder(user_input)
-    return text, image
-
-btn.click(fn=on_click, inputs=[chatbot], outputs=[output_text, output_image])
-
-demo.launch()
-
+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()