Amsterdam / app.py
Hodely's picture
Update app.py
02ce9bf verified
raw
history blame
1.34 kB
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()