go back

app.py

@@ -1,8 +1,6 @@
 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
@@ -170,40 +168,39 @@ 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)
-
-
-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,
-    )
+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]
     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_first_stage":"COCO-Stuff", "faceshq":"FacesHQ", "sflckr":"S-FLCKR"}
+                    "wikiart_1024":"WikiArt 1024", "wikiart_16384":"WikiArt 16384", "coco":"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 +342,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, model_name):
+def throttled_inference(text, seed, step_size, max_iterations, width, height):
     global inferences_running
     current = inferences_running
     if current >= 3:
@@ -354,7 +351,7 @@ def throttled_inference(text, seed, step_size, max_iterations, width, height, mo
     print(f"Inference starting when we already had {current} running")
     inferences_running += 1
     try:
-        return inference(text, seed, step_size, max_iterations, width, height, model_name)
+        return inference(text, seed, step_size, max_iterations, width, height)
     finally:
         print("Inference finished")
         inferences_running -= 1
@@ -369,15 +366,14 @@ 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, "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"]
+              ['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]
     ]
     ).launch(debug=True)