Update app.py

app.py

@@ -40,19 +40,22 @@ 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()))
-    
+    #display(Image.fromarray(reshaped.numpy()))
+    #Image.fromarray(reshaped.numpy()).save('image.png')
+
+
 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())
-        
+    img = Image.fromarray(reshaped.numpy())
+    #img.save('img.png')
+    return img
+
 # Sampling parameters
 batch_size = 1
 guidance_scale = 3.0
@@ -122,7 +125,7 @@ def run(prompt):
 
     # Show the output
     print('    # Show the output')
-    show_images(samples)
+    #show_images(samples)
     ##############################
     # Upsample the 64x64 samples #
     ##############################
@@ -167,7 +170,6 @@ def run(prompt):
     
     # Show the output
     print('# Show the output')
-    #    show_images(up_samples)
     out_images = get_images(up_samples)
     
     return out_images