Update app.py

app.py

@@ -17,7 +17,7 @@ model = PaletteModelV2(c_in=2, c_out=1, num_classes=5, image_size=256, true_img_
 ckpt = torch.load('ema_ckpt_cond.pt')
 model.load_state_dict(ckpt)
 
-diffusion = Diffusion_cond(noise_steps=1000, img_size=256, device=device)
+diffusion = Diffusion_cond(img_size=256, device=device)
 model.eval()
 
 transform_hmi = transforms.Compose([
@@ -30,8 +30,13 @@ transform_hmi = transforms.Compose([
 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())
-    return generated_image_pil
+    # 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 = np.squeeze(img.cpu().numpy())
+    v = Image.fromarray(img)                # Create a PIL Image from array
+    v = v.transpose(Image.FLIP_TOP_BOTTOM)  
+    
+    return v
 
 # Create Gradio interface
 iface = gr.Interface(