version 0.1 commit

app.py

@@ -10,13 +10,14 @@ from image_generator import generate_images
 
 def image_generation(model, number_of_images=1):
     G = MyGenerator.from_pretrained("Cropinky/projected_gan_impressionism")
-    generate_images()
-    return f"generating {number_of_images} images from {model}"
+    img = generate_images(model)
+    #return f"generating {number_of_images} images from {model}"
+    return img
 if __name__ == "__main__":
 
-    inputs = gr.inputs.Radio(["Abstract Expressionism", "Impressionism", "Cubism", "Minimalism", "Pop Art", "Color Field", "Hana Hanak houses"])
-    #outputs = gr.outputs.Image(label="Output Image")
-    outputs = "text"
+    inputs = gr.inputs.Radio(["Abstract Expressionism", "Impressionism", "Cubism", "Pop Art", "Color Field", "Hana Hanak houses"])
+    outputs = gr.outputs.Image(label="Generated Image", type="pil")
+    #outputs = "text"
     title = "Projected GAN for painting generation"
     description = "Choose your artistic direction "
     article = "<p style='text-align: center'><a href='https://github.com/autonomousvision/projected_gan'>Official projected GAN github repo + paper</a></p>"

image_generator.py

@@ -8,15 +8,16 @@
 
 """Generate images using pretrained network pickle."""
 
+from ast import parse
 import os
+from pyexpat import model
 import re
 from typing import List, Optional, Tuple, Union
-import click
 import numpy as np
 import PIL.Image
 import torch
 from networks_fastgan import MyGenerator
-
+import random
 #----------------------------------------------------------------------------
 
 def parse_range(s: Union[str, List]) -> List[int]:
@@ -65,22 +66,14 @@ def make_transform(translate: Tuple[float,float], angle: float):
 
 #----------------------------------------------------------------------------
 
-@click.command()
-@click.option('--seeds', type=parse_range, help='List of random seeds (e.g., \'0,1,4-6\')', default="10-11", required=True)
-@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
-@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)')
-@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)
-@click.option('--translate', help='Translate XY-coordinate (e.g. \'0.3,1\')', type=parse_vec2, default='0,0', show_default=True, metavar='VEC2')
-@click.option('--rotate', help='Rotation angle in degrees', type=float, default=0, show_default=True, metavar='ANGLE')
-@click.option('--outdir', help='Where to save the output images', type=str, default="out", required=True, metavar='DIR')
 def generate_images(
-    seeds: List[int],
-    truncation_psi: float,
-    noise_mode: str,
-    outdir: str,
-    translate: Tuple[float,float],
-    rotate: float,
-    class_idx: Optional[int]
+    model_path, 
+    seeds = "10-12",
+    truncation_psi = 1.0,
+    noise_mode = "const",
+    outdir = "out",
+    translate = "0,0",
+    rotate = 0,
 ):
     """Generate images using pretrained network pickle.
 
@@ -96,13 +89,30 @@ def generate_images(
     python gen_images.py --outdir=out --trunc=0.7 --seeds=600-605 \\
         --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-metfacesu-1024x1024.pkl
     """
-
+    model_owner = "Cropinky"
+    #inputs = gr.inputs.Radio(["Abstract Expressionism", "Impressionism", "Cubism", "Minimalism", "Pop Art", "Color Field", "Hana Hanak houses"])
+    model_path_dict = {
+        'Impressionism' : 'projected_gan_impressionism',
+        'Cubism' : 'projected_gan_cubism', 
+        'Abstract Expressionism' : 'projected_gan_abstract_expressionism', 
+        'Pop Art' : 'projected_gan_pop_art',
+        'Minimalism' : 'projected_gan_minimalism',
+        'Color Field' : 'projected_gan_color_field',
+        'Hana Hanak houses' : 'projected_gana_hana' 
+    }
+    
+    model_path = model_owner + "/" + model_path_dict[model_path]
+    print(model_path)
+    seeds = parse_range(seeds)
+    print(seeds)
+    seeds=[random.randint(1,99)]
     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
-    G = MyGenerator.from_pretrained("Cropinky/projected_gan_impressionism")
+    G = MyGenerator.from_pretrained(model_path)
     os.makedirs(outdir, exist_ok=True)
 
     # Labels.
     label = torch.zeros([1, G.c_dim], device=device)
+    """
     if G.c_dim != 0:
         if class_idx is None:
             raise click.ClickException('Must specify class label with --class when using a conditional network')
@@ -110,6 +120,7 @@ def generate_images(
     else:
         if class_idx is not None:
             print ('warn: --class=lbl ignored when running on an unconditional network')
+    """
 
     # Generate images.
     for seed_idx, seed in enumerate(seeds):
@@ -125,7 +136,9 @@ def generate_images(
 
         img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode)
         img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-        PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png')
+        img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
+        #PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png')
+    return img
 
 
 #----------------------------------------------------------------------------