Update app.py

app.py

@@ -31,7 +31,7 @@ def generate_image(seed_image):
     seed_image_tensor = transform_hmi(Image.open(seed_image)).reshape(1, 1, 256, 256).to(device)
     generated_image = diffusion.sample(model, y=seed_image_tensor, labels=None, n=1)
     # generated_image_pil = transforms.ToPILImage()(generated_image.squeeze().cpu())
-    img = generated_image[0].permute(1, 2, 0)              # Permute dimensions to height x width x channels
+    img = generated_image[0].reshape(1, 256, 256).permute(1, 2, 0)              # Permute dimensions to height x width x channels
     img = np.squeeze(img.cpu().numpy())
     v = Image.fromarray(img)                # Create a PIL Image from array
     v = v.transpose(Image.FLIP_TOP_BOTTOM)