import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from diffusers import StableDiffusionPipeline
import torch

# Configuración de Llama o Nous para texto
model_name = "meta-llama/Llama-2-7b-chat-hf"  # o el modelo Nous que prefieras
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Configuración de Stable Diffusion para imágenes
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to("cuda")

# Función de generación de texto
def generate_text(prompt):
    response = generator(prompt, max_length=60, num_return_sequences=1, temperature=0.5, top_p=0.85)
    return response[0]['generated_text']

# Función de generación de imagen
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Crear la interfaz de Gradio
iface_text = gr.Interface(fn=generate_text, inputs="text", outputs="text", description="Generador de Texto con Llama/Nous")
iface_image = gr.Interface(fn=generate_image, inputs="text", outputs="image", description="Generador de Imágenes con Stable Diffusion")

# Ejecutar ambas interfaces juntas
app = gr.TabbedInterface([iface_text, iface_image], ["Texto", "Imagen"])
app.launch()