Update app.py

app.py

@@ -39,11 +39,20 @@ if has_cuda:
 model_up.to(device)
 model_up.load_state_dict(load_checkpoint('upsample', device))
 print('total upsampler parameters', sum(x.numel() for x in model_up.parameters()))
+
+
 def show_images(batch: th.Tensor):
     """ Display a batch of images inline. """
     scaled = ((batch + 1)*127.5).round().clamp(0,255).to(th.uint8).cpu()
     reshaped = scaled.permute(2, 0, 3, 1).reshape([batch.shape[2], -1, 3])
     display(Image.fromarray(reshaped.numpy()))
+    
+def get_images(batch: th.Tensor):
+    """ Display a batch of images inline. """
+    scaled = ((batch + 1)*127.5).round().clamp(0,255).to(th.uint8).cpu()
+    reshaped = scaled.permute(2, 0, 3, 1).reshape([batch.shape[2], -1, 3])
+    return Image.fromarray(reshaped.numpy())
+        
 # Sampling parameters
 batch_size = 1
 guidance_scale = 3.0
@@ -159,8 +168,9 @@ def run(prompt):
     # Show the output
     print('# Show the output')
     #    show_images(up_samples)
+    out_images = get_images(up_samples)
     
-    return up_samples
+    return out_images