simon0820 commited on
Commit
59de3af
1 Parent(s): 5c09100

reparando errores

Browse files
Files changed (2) hide show
  1. app.py +51 -26
  2. utils.py +10 -6
app.py CHANGED
@@ -1,34 +1,59 @@
1
  import streamlit as st
2
 
3
- from utils import load_model,generate
4
-
5
- st.title("butterfly_generator")
6
- st.write("this is a ligth GAN generator model")
7
-
8
- st.sidebar.subheader("this butterfly is created by IA")
9
- st.sidebar.image("assets/logo.png",width=200)
10
-
11
- model_id="ceyda/butterfly_cropped_uniq1K_512"
12
- model_gan=load_model(model_id)
13
- n_fly=4
14
- def run():
15
- with st.spinner("generating"):
16
- ims=generate(model_gan,n_fly)
17
- st.session_state["ims"]=ims
18
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
  if "ims" not in st.session_state:
20
- st.session_state["ims"]=None
21
- run()
22
 
23
- ims=st.session_state["ims"]
 
24
 
25
- run_button=st.button(
26
- "click me",
27
- on_click=run
28
- )
 
 
29
 
30
  if ims is not None:
31
- cols=st.columns(n_fly)
32
  for j, im in enumerate(ims):
33
- i=j%n_fly
34
- cols[i].image(im,use_column_width=True)
 
1
  import streamlit as st
2
 
3
+ from utils import carga_modelo, genera
4
+
5
+ ## Página principal
6
+ st.title("Butterfly GAN (GAN de mariposas)")
7
+ st.write(
8
+ "Modelo Light-GAN entrenado con 1000 imágenes de mariposas tomadas de la colección del Museo Smithsonian."
9
+ )
10
+
11
+ ## Barra lateral
12
+ st.sidebar.subheader("¡Esta mariposa no existe! Ni en América Latina 🤯.")
13
+ st.sidebar.image("assets/logo.png", width=200)
14
+ st.sidebar.caption(
15
+ f"[Modelo](https://huggingface.co/ceyda/butterfly_cropped_uniq1K_512) y [Dataset](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) usados."
16
+ )
17
+ st.sidebar.caption(f"*Disclaimers:*")
18
+ st.sidebar.caption(
19
+ "* Este demo es una versión simplificada del creado por [Ceyda Cinarel](https://github.com/cceyda) y [Jonathan Whitaker](https://datasciencecastnet.home.blog/) ([link](https://huggingface.co/spaces/huggan/butterfly-gan)) durante el hackathon [HugGan](https://github.com/huggingface/community-events). Cualquier error se atribuye a [Omar Espejel](https://twitter.com/espejelomar)."
20
+ )
21
+ st.sidebar.caption(
22
+ "* Modelo basado en el [paper](https://openreview.net/forum?id=1Fqg133qRaI) *Towards Faster and Stabilized GAN Training for High-fidelity Few-shot Image Synthesis*."
23
+ )
24
+
25
+ ## Cargamos modelo
26
+ repo_id = "ceyda/butterfly_cropped_uniq1K_512"
27
+ version_modelo = "57d36a15546909557d9f967f47713236c8288838"
28
+ modelo_gan = carga_modelo(repo_id, version_modelo)
29
+
30
+ ## Generamos 4 mariposas
31
+ n_mariposas = 4
32
+
33
+ ## Función que genera mariposas y lo guarda como un estado de la sesión
34
+ def corre():
35
+ with st.spinner("Generando, espera un poco..."):
36
+ ims = genera(modelo_gan, n_mariposas)
37
+ st.session_state["ims"] = ims
38
+
39
+
40
+ ## Si no hay una imagen generada entonces generala
41
  if "ims" not in st.session_state:
42
+ st.session_state["ims"] = None
43
+ corre()
44
 
45
+ ## ims contiene las imágenes generadas
46
+ ims = st.session_state["ims"]
47
 
48
+ ## Si la usuaria da click en el botón entonces corremos la función genera()
49
+ corre_boton = st.button(
50
+ "Genera mariposas, porfa.",
51
+ on_click=corre,
52
+ help="Estamos en pleno vuelo, puede tardar.",
53
+ )
54
 
55
  if ims is not None:
56
+ cols = st.columns(n_mariposas)
57
  for j, im in enumerate(ims):
58
+ i = j % n_mariposas
59
+ cols[i].image(im, use_column_width=True)
utils.py CHANGED
@@ -2,13 +2,17 @@ import numpy as np
2
  import torch
3
  from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
4
 
5
- def load_model(model_name="ceyda/butterfly_cropped_uniq1K_512", model_version=None):
6
- gan=LightweightGAN.from_pretrained(model_name,version=model_version)
 
 
7
  gan.eval()
8
  return gan
9
 
10
- def generate(gan,batch_size=1):
 
 
11
  with torch.no_grad():
12
- ims=gan.G(torch.randn(batch_size,gan.latent_dim)).clamp(0.0,1.0)*255
13
- ims=ims.permute(0,2,3,1).deatch().cpu().numpy().astype(np.uint8)
14
- return ims
 
2
  import torch
3
  from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
4
 
5
+
6
+ ## Cargamos el modelo desde el Hub de Hugging Face
7
+ def carga_modelo(model_name="ceyda/butterfly_cropped_uniq1K_512", model_version=None):
8
+ gan = LightweightGAN.from_pretrained(model_name, version=model_version)
9
  gan.eval()
10
  return gan
11
 
12
+
13
+ ## Usamos el modelo GAN para generar imágenes
14
+ def genera(gan, batch_size=1):
15
  with torch.no_grad():
16
+ ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0.0, 1.0) * 255
17
+ ims = ims.permute(0, 2, 3, 1).detach().cpu().numpy().astype(np.uint8)
18
+ return ims