Update app.py

app.py

@@ -6,7 +6,7 @@ import math
 from pathlib import Path
 import sys
 sys.path.insert(1, './taming-transformers')
-#from IPython import display
+# from IPython import display
 from base64 import b64encode
 from omegaconf import OmegaConf
 from PIL import Image
@@ -25,6 +25,7 @@ from PIL import ImageFile, Image
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 import gradio as gr
 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([]))
 def lanczos(x, a):
@@ -189,7 +190,8 @@ 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):
+def inference(text, seed, step_size, max_iterations, width, height):
+    size=[width, height]
     texts = text
     target_images = ""
     max_iterations = max_iterations
@@ -221,7 +223,7 @@ def inference(text, seed, step_size, max_iterations):
     cut_size = perceptor.visual.input_resolution
     f = 2**(model.decoder.num_resolutions - 1)
     make_cutouts = MakeCutouts(cut_size, args.cutn, cut_pow=args.cut_pow)
-    toksX, toksY = args.size[0] // f, args.size[1] // f
+    toksX, toksY = size[0] // f, size[1] // f
     sideX, sideY = toksX * f, toksY * f
     if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':
         e_dim = 256
@@ -237,11 +239,11 @@ def inference(text, seed, step_size, max_iterations):
     # z_max = model.quantize.embedding.weight.max(dim=0).values[None, :, None, None]
     # normalize_imagenet = transforms.Normalize(mean=[0.485, 0.456, 0.406],
     #                                            std=[0.229, 0.224, 0.225])
-    if args.init_image:
-        if 'http' in args.init_image:
-            img = Image.open(urlopen(args.init_image))
+    if init_image:
+        if 'http' in init_image:
+            img = Image.open(urlopen(init_image))
         else:
-            img = Image.open(args.init_image)
+            img = Image.open(init_image)
         pil_image = img.convert('RGB')
         pil_image = pil_image.resize((sideX, sideY), Image.LANCZOS)
         pil_tensor = TF.to_tensor(pil_image)
@@ -288,8 +290,8 @@ def inference(text, seed, step_size, max_iterations):
         losses_str = ', '.join(f'{loss.item():g}' for loss in losses)
         tqdm.write(f'i: {i}, loss: {sum(losses).item():g}, losses: {losses_str}')
         out = synth(z)
-        #TF.to_pil_image(out[0].cpu()).save('progress.png')
-        #display.display(display.Image('progress.png'))
+        # TF.to_pil_image(out[0].cpu()).save('progress.png')
+        # display.display(display.Image('progress.png'))
     def ascend_txt():
         # global i
         out = synth(z)
@@ -335,7 +337,7 @@ def load_image( infilename ) :
     img.load()
     data = np.asarray( img, dtype="int32" )
     return data
-def throttled_inference(text, seed, step_size, max_iterations):
+def throttled_inference(text, seed, step_size, max_iterations, width, height):
     global inferences_running
     current = inferences_running
     if current >= 2:
@@ -344,7 +346,7 @@ def throttled_inference(text, seed, step_size, max_iterations):
     print(f"Inference starting when we already had {current} running")
     inferences_running += 1
     try:
-        return inference(text, seed, step_size, max_iterations)
+        return inference(text, seed, step_size, max_iterations, width, height)
     finally:
         print("Inference finished")
         inferences_running -= 1
@@ -357,14 +359,16 @@ gr.Interface(
      gr.inputs.Number(default=42, label="seed"),
      gr.inputs.Slider(minimum=0.1, maximum=0.9, default=0.23, label='step size'),
     gr.inputs.Slider(minimum=100, maximum=150, default=100, 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.outputs.Image(type="numpy", label="Output"),
     title=title,
     description=description,
     article=article,
     examples=[
-              ['a garden by james gurney',42,0.23, 100],
-              ['coral reef city artstationHQ',1000,0.6, 110],
-              ['a cabin in the mountains unreal engine',98,0.3, 120]
+              ['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)