Add conditioned image

app.py

@@ -35,6 +35,10 @@ 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
+    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)
+    
     # generate (no beam search)
     context = np.concatenate((np.full((batch_size, 1), model.config.vocab_size - 1), primers), axis=1)
     context = torch.tensor(context).to(device)
@@ -52,16 +56,17 @@ def process_image(image):
     # return as PIL Image
     completion = Image.fromarray(result)
 
-    return completion
+    return [primers_img, 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"),
+                     outputs=[gr.outputs.Image(type="pil", label=labels[idx]) for idx in range(2)],
                      title=title,
                      description=description,
                      article=article,