Added light and dark models

app.py

@@ -465,12 +465,26 @@ class GeneratorWithPyTorchModelHubMixin(gan_new.__class__, PyTorchModelHubMixin)
 gan_new.__class__ = GeneratorWithPyTorchModelHubMixin
 gan_new = gan_new.from_pretrained('johnowhitaker/orbgan_e1', latent_dim=256, image_size=256, attn_res_layers = [32])
 
-def gen_ims(n_rows):
-  ims = gan_new(torch.randn(int(n_rows)**2, 256)).clamp_(0., 1.)
+gan_light = Generator(latent_dim=256, image_size=256, attn_res_layers = [32])
+gan_light.__class__ = GeneratorWithPyTorchModelHubMixin
+gan_light = gan_light.from_pretrained('johnowhitaker/orbgan_ligt', latent_dim=256, image_size=256, attn_res_layers = [32])
+
+gan_dark = Generator(latent_dim=256, image_size=256, attn_res_layers = [32])
+gan_dark.__class__ = GeneratorWithPyTorchModelHubMixin
+gan_dark = gan_dark.from_pretrained('johnowhitaker/orbgan_ligt', latent_dim=256, image_size=256, attn_res_layers = [32])
+
+def gen_ims(n_rows, model='both'):
+  if model == "both":
+    ims = gan_new(torch.randn(int(n_rows)**2, 256)).clamp_(0., 1.)
+  if model == "light":
+    ims = gan_light(torch.randn(int(n_rows)**2, 256)).clamp_(0., 1.)
+  if model == "dark":
+    ims = gan_dark(torch.randn(int(n_rows)**2, 256)).clamp_(0., 1.)
   grid = torchvision.utils.make_grid(ims, nrow=int(n_rows)).permute(1, 2, 0).detach().cpu().numpy()
   return (grid*255).astype(np.uint8)
 iface = gr.Interface(fn=gen_ims, 
-  inputs=[gr.inputs.Slider(minimum=1, maximum=6, step=1, default=3,label="N rows")], 
+  inputs=[gr.inputs.Slider(minimum=1, maximum=6, step=1, default=3,label="N rows"),
+  gradio.inputs.Dropdown(["both", "light", "dark"], type="value", default="dark", label="Orb Type (model)", optional=False)], 
   outputs=[gr.outputs.Image(type="numpy", label="Generated Images")],
   title='Demo for https://huggingface.co/johnowhitaker/orbgan_e1'
 )