remove post-processing

app.py

@@ -25,23 +25,15 @@ def predict(steps=100, seed=42,scheduler="ddim"):
     torch.cuda.empty_cache()
     generator = torch.manual_seed(seed)
     if(scheduler == "ddim"):
-        image = ddim_pipeline(generator=generator, num_inference_steps=steps)["sample"]
+        images = ddim_pipeline(generator=generator, num_inference_steps=steps)["sample"]
     elif(scheduler == "ddpm"):
-        image = ddpm_pipeline(generator=generator)["sample"]
+        images = ddpm_pipeline(generator=generator)["sample"]
     elif(scheduler == "pndm"):
         if(steps > 100):
             steps = 100
-        image = pndm_pipeline(generator=generator, num_inference_steps=steps)["sample"]
-
-    image_processed = image.cpu().permute(0, 2, 3, 1)
-    if scheduler == "pndm":
-        image_processed = (image_processed + 1.0) / 2
-        image_processed = torch.clamp(image_processed, 0.0, 1.0)
-        image_processed = image_processed * 255
-    else:
-        image_processed = (image_processed + 1.0) * 127.5
-    image_processed = image_processed.detach().numpy().astype(np.uint8)
-    return(PIL.Image.fromarray(image_processed[0]))
+        images = pndm_pipeline(generator=generator, num_inference_steps=steps)["sample"]
+
+    return(images[0])
 
 
 random_seed = random.randint(0, 2147483647)