Update app.py

app.py

@@ -24,8 +24,8 @@ for idx, url in enumerate(urls):
   image.save(f"image_{idx}.png")
 
 def process_image(image):
-    # prepare 8 images, shape (8, 1024)
-    batch_size = 8
+    # prepare 7 images, shape (7, 1024)
+    batch_size = 7
     encoding = feature_extractor([image for _ in range(batch_size)], return_tensors="pt")
 
     # create primers
@@ -35,9 +35,9 @@ def process_image(image):
     n_px_crop = 16
     primers = samples.reshape(-1,n_px*n_px)[:,:n_px_crop*n_px] # crop top n_px_crop rows. These will be the conditioning tokens
     
-    # get conditioned image (from first primer tensor) by converting color clusters back to pixels
+    # get conditioned image (from first primer tensor), padded with black pixels to be 32x32
     primers_img = np.reshape(np.rint(127.5 * (clusters[primers[0]] + 1.0)), [n_px_crop,n_px, 3]).astype(np.uint8) 
-    primers_img = Image.fromarray(primers_img)
+    primers_img = np.pad(primers_img, pad_width=((0,16), (0,0), (0,0)), mode="constant")
     
     # generate (no beam search)
     context = np.concatenate((np.full((batch_size, 1), model.config.vocab_size - 1), primers), axis=1)
@@ -48,6 +48,8 @@ def process_image(image):
     samples = output[:,1:].cpu().detach().numpy()
     samples_img = [np.reshape(np.rint(127.5 * (clusters[s] + 1.0)), [n_px, n_px, 3]).astype(np.uint8) for s in samples] 
     
+    samples_img = [primers_img] + samples_img
+    
     # stack images horizontally
     row1 = np.hstack(samples_img[:4])
     row2 = np.hstack(samples_img[4:])
@@ -56,17 +58,16 @@ def process_image(image):
     # return as PIL Image
     completion = Image.fromarray(result)
 
-    return [primers_img, completion]
+    return completion
 
 title = "Interactive demo: ImageGPT"
 description = "Demo for OpenAI's ImageGPT: Generative Pretraining from Pixels. To use it, simply upload an image or use the example image below and click 'submit'. Results will show up in a few seconds."
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2109.10282'>ImageGPT: Generative Pretraining from Pixels</a> | <a href='https://openai.com/blog/image-gpt/'>Official blog</a></p>"
 examples =[f"image_{idx}.png" for idx in range(len(urls))]
 
-labels = ["Conditioned image:", "Completions:"]
 iface = gr.Interface(fn=process_image, 
                      inputs=gr.inputs.Image(type="pil"), 
-                     outputs=[gr.outputs.Image(type="pil", label=labels[idx]) for idx in range(2)],
+                     outputs=gr.outputs.Image(type="pil", label="Model input + completions"),
                      title=title,
                      description=description,
                      article=article,