juanelot commited on
Commit
b46cce9
verified
1 Parent(s): 91703bb
Files changed (1) hide show
  1. app.py +9 -23
app.py CHANGED
@@ -1,33 +1,19 @@
1
  import gradio as gr
2
- from diffusers import DiffusionPipeline
3
  import torch
4
- from transformers.utils import move_cache
5
- import logging
6
 
7
- # Suprimir advertencias de FutureWarning
8
- logging.getLogger("transformers.configuration_utils").setLevel(logging.ERROR)
9
- logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
10
-
11
- # Mover la cach茅 si es necesario
12
- move_cache()
13
-
14
- # Cargar el modelo base y los pesos LoRA
15
- model_name = "stabilityai/stable-diffusion-xl-base-1.0"
16
- lora_weights = "ZB-Tech/Text-to-Image"
17
- pipeline = DiffusionPipeline.from_pretrained(model_name, lora_backend="cuda")
18
- pipeline.load_lora_weights(lora_weights)
19
-
20
- # Mover el modelo a la GPU si est谩 disponible
21
- device = "cuda" if torch.cuda.is_available() else "cpu"
22
- pipeline = pipeline.to(device)
23
 
24
  # Definir la funci贸n de generaci贸n de imagen
25
  def generate_image(prompt):
26
- with torch.autocast(device):
27
- image = pipeline(prompt).images[0]
28
- return image
 
29
 
30
- # Crear la interfaz de Gradio con CSS personalizado
31
  iface = gr.Interface(
32
  fn=generate_image,
33
  inputs=gr.Textbox(lines=5, label="Descripci贸n de la imagen", placeholder="Introduce el texto aqu铆..."),
 
1
  import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForTextToImageGeneration
3
  import torch
 
 
4
 
5
+ # Cargar el tokenizador y el modelo
6
+ tokenizer = AutoTokenizer.from_pretrained("ZB-Tech/Text-to-Image")
7
+ model = AutoModelForTextToImageGeneration.from_pretrained("ZB-Tech/Text-to-Image")
 
 
 
 
 
 
 
 
 
 
 
 
 
8
 
9
  # Definir la funci贸n de generaci贸n de imagen
10
  def generate_image(prompt):
11
+ inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
12
+ with torch.no_grad():
13
+ output = model.generate(**inputs)
14
+ return output
15
 
16
+ # Crear la interfaz de Gradio
17
  iface = gr.Interface(
18
  fn=generate_image,
19
  inputs=gr.Textbox(lines=5, label="Descripci贸n de la imagen", placeholder="Introduce el texto aqu铆..."),