Update app.py

app.py

@@ -1,6 +1,8 @@
 import torch
 torch.hub.download_url_to_file('http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_16384.yaml', 'vqgan_imagenet_f16_16384.yaml')
 torch.hub.download_url_to_file('http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_16384.ckpt', 'vqgan_imagenet_f16_16384.ckpt')
+torch.hub.download_url_to_file('http://batbot.tv/misc/coco_first_stage.yaml', 'coco_first_stage.yaml')
+torch.hub.download_url_to_file('http://batbot.tv/misc/coco_first_stage.ckpt', 'coco_first_stage.ckpt')
 import argparse
 import math
 from pathlib import Path
@@ -25,12 +27,9 @@ from PIL import ImageFile, Image
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 import gradio as gr
 import nvidia_smi
-
 nvidia_smi.nvmlInit()
 handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
 # card id 0 hardcoded here, there is also a call to get all available card ids, so we could iterate
-
-
 torch.hub.download_url_to_file('https://i.imgur.com/WEHmKef.jpg', 'gpu.jpg')
 def sinc(x):
     return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))
@@ -171,39 +170,40 @@ def resize_image(image, out_size):
     area = min(image.size[0] * image.size[1], out_size[0] * out_size[1])
     size = round((area * ratio)**0.5), round((area / ratio)**0.5)
     return image.resize(size, Image.LANCZOS)
-model_name = "vqgan_imagenet_f16_16384" 
-images_interval =  50
-width =  280
-height = 280
-init_image = ""
-seed = 42
-args = argparse.Namespace(
-    noise_prompt_seeds=[],
-    noise_prompt_weights=[],
-    size=[width, height],
-    init_image=init_image,
-    init_weight=0.,
-    clip_model='ViT-B/32',
-    vqgan_config=f'{model_name}.yaml',
-    vqgan_checkpoint=f'{model_name}.ckpt',
-    step_size=0.15,
-    cutn=4,
-    cut_pow=1.,
-    display_freq=images_interval,
-    seed=seed,
-)
-device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-print('Using device:', device)
-model = load_vqgan_model(args.vqgan_config, args.vqgan_checkpoint).to(device)
-perceptor = clip.load(args.clip_model, jit=False)[0].eval().requires_grad_(False).to(device)
-def inference(text, seed, step_size, max_iterations, width, height):
-    size=[width, height]
+
+
+def inference(text, seed, step_size, max_iterations, width, height, model_name):
+    args = argparse.Namespace(
+        noise_prompt_seeds=[],
+        noise_prompt_weights=[],
+        size=[width, height],
+        init_image="",
+        init_weight=0.,
+        clip_model='ViT-B/32',
+        vqgan_config=f'{model_name}.yaml',
+        vqgan_checkpoint=f'{model_name}.ckpt',
+        step_size=step_size,
+        cutn=4,
+        cut_pow=1.,
+        display_freq=50,
+        seed=seed,
+    )
     texts = text
     target_images = ""
     max_iterations = max_iterations
+    model_name = model_name
     model_names={"vqgan_imagenet_f16_16384": 'ImageNet 16384',"vqgan_imagenet_f16_1024":"ImageNet 1024", 'vqgan_openimages_f16_8192':'OpenImages 8912',
-                    "wikiart_1024":"WikiArt 1024", "wikiart_16384":"WikiArt 16384", "coco":"COCO-Stuff", "faceshq":"FacesHQ", "sflckr":"S-FLCKR"}
+                    "wikiart_1024":"WikiArt 1024", "wikiart_16384":"WikiArt 16384", "coco_first_stage":"COCO-Stuff", "faceshq":"FacesHQ", "sflckr":"S-FLCKR"}
     name_model = model_names[model_name]
+    init_image = ""
+    size=[width, height]
+    seed=seed
+    step_size=step_size
+   
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    print('Using device:', device)
+    model = load_vqgan_model(args.vqgan_config, args.vqgan_checkpoint).to(device)
+    perceptor = clip.load(args.clip_model, jit=False)[0].eval().requires_grad_(False).to(device)
     if target_images == "None" or not target_images:
         target_images = []
     else:
@@ -345,7 +345,7 @@ def load_image( infilename ) :
     img.load()
     data = np.asarray( img, dtype="int32" )
     return data
-def throttled_inference(text, seed, step_size, max_iterations, width, height):
+def throttled_inference(text, seed, step_size, max_iterations, width, height, model_name):
     global inferences_running
     current = inferences_running
     if current >= 3:
@@ -354,7 +354,7 @@ def throttled_inference(text, seed, step_size, max_iterations, width, height):
     print(f"Inference starting when we already had {current} running")
     inferences_running += 1
     try:
-        return inference(text, seed, step_size, max_iterations, width, height)
+        return inference(text, seed, step_size, max_iterations, width, height, model_name)
     finally:
         print("Inference finished")
         inferences_running -= 1
@@ -369,14 +369,15 @@ gr.Interface(
     gr.inputs.Slider(minimum=25, maximum=150, default=80, label='max iterations', step=1),
     gr.inputs.Slider(minimum=200, maximum=280, default=256, label='width', step=1),
     gr.inputs.Slider(minimum=200, maximum=280, default=256, label='height', step=1),
+    gr.inputs.Dropdown(choices=["vqgan_imagenet_f16_16384", "coco_first_stage"], type="value", default="vqgan_imagenet_f16_16384", label="Model Name")
      ], 
     gr.outputs.Image(type="numpy", label="Output"),
     title=title,
     description=description,
     article=article,
     examples=[
-              ['a garden by james gurney',42,0.16, 100, 256, 256],
-              ['coral reef city artstationHQ',1000,0.6, 110, 200, 200],
-              ['a cabin in the mountains unreal engine',98,0.3, 120, 280, 280]
+              ['a garden by james gurney',42,0.16, 100, 256, 256, "vqgan_imagenet_f16_16384"],
+              ['coral reef city artstationHQ',1000,0.6, 110, 200, 200, "vqgan_imagenet_f16_16384"],
+              ['a cabin in the mountains unreal engine',98,0.3, 120, 280, 280, "vqgan_imagenet_f16_16384"]
     ]
     ).launch(debug=True)