Stack images horizontally

app.py

@@ -37,19 +37,17 @@ def process_image(image):
     context = torch.tensor(context).to(device)
     output = model.generate(input_ids=context, max_length=n_px*n_px + 1, temperature=1.0, do_sample=True, top_k=40)
 
-    # decode back to images
+    # decode back to images (convert color cluster tokens back to pixels)
     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] # convert color cluster tokens back to pixels
+    samples_img = [np.reshape(np.rint(127.5 * (clusters[s] + 1.0)), [n_px, n_px, 3]).astype(np.uint8) for s in samples] 
     
-    # save as list of files
-    completions = []
-    output_dir = '.'
-    for i in range(len(samples_img)):
-      fname = os.path.join(output_dir, "completion" + str(i) + ".png")
-      plt.imsave(fname=fname, arr=samples_img[i], format='png')
-      completions.append(fname)
+    # stack images horizontally
+    result = np.hstack(samples_img)
+    
+    # return as PIL Image
+    completion = Image.fromarray(result)
 
-    return completions
+    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."
@@ -58,7 +56,7 @@ examples =[["image_0.png"]]
 
 iface = gr.Interface(fn=process_image, 
                      inputs=gr.inputs.Image(type="pil"), 
-                     outputs=[gr.outputs.Image(type='file', label=f'completion_{i}') for i in range(8)],
+                     outputs=gr.outputs.Image(type="pil"),
                      title=title,
                      description=description,
                      article=article,