Update app.py

app.py

@@ -24,9 +24,7 @@ import imageio
 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):
@@ -166,8 +164,6 @@ 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
@@ -193,11 +189,10 @@ 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):
+def inference(text, seed, step_size, max_iterations):
     texts = text
     target_images = ""
-    max_iterations = 100
+    max_iterations = max_iterations
     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"}
     name_model = model_names[model_name]
@@ -334,31 +329,25 @@ def inference(text, seed, step_size):
     except KeyboardInterrupt:
         pass
     return image
-
 inferences_running = 0
-
 def load_image( infilename ) :
     img = Image.open( infilename )
     img.load()
     data = np.asarray( img, dtype="int32" )
     return data
-
-def throttled_inference(text, seed, step_size):
+def throttled_inference(text, seed, step_size, max_iterations):
     global inferences_running
     current = inferences_running
     if current >= 2:
         print(f"Rejected inference when we already had {current} running")
         return load_image("./gpu.jpg")
-
     print(f"Inference starting when we already had {current} running")
     inferences_running += 1
     try:
-        return inference(text, seed, step_size)
+        return inference(text, seed, step_size, max_iterations)
     finally:
         print("Inference finished")
         inferences_running -= 1
-
-
 title = "VQGAN + CLIP"
 description = "Gradio demo for VQGAN + CLIP. To use it, simply add your text, or click one of the examples to load them. Read more at the links below. Please click submit only once. Results will show up in under a minute."
 article = "<p style='text-align: center'>Originally made by Katherine Crowson (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings). The original BigGAN+CLIP method was by https://twitter.com/advadnoun. Added some explanations and modifications by Eleiber#8347, pooling trick by Crimeacs#8222 (https://twitter.com/EarthML1) and the GUI was made with the help of Abulafia#3734. | <a href='https://colab.research.google.com/drive/1ZAus_gn2RhTZWzOWUpPERNC0Q8OhZRTZ'>Colab</a> | <a href='https://github.com/CompVis/taming-transformers'>Taming Transformers Github Repo</a> | <a href='https://github.com/openai/CLIP'>CLIP Github Repo</a> | Special thanks to BoneAmputee (https://twitter.com/BoneAmputee) for suggestions and advice</p>"
@@ -366,15 +355,16 @@ gr.Interface(
     throttled_inference, 
     [gr.inputs.Textbox(label="Input"),
      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=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.outputs.Image(type="numpy", label="Output"),
     title=title,
     description=description,
     article=article,
     examples=[
-              ['a garden by james gurney',42,0.23],
-              ['coral reef city artstationHQ',1000,0.6],
-              ['a cabin in the mountains unreal engine',98,0.3]
+              ['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]
     ]
     ).launch(debug=True)