Add sample and bring back the steps slider

app.py

@@ -21,18 +21,15 @@ ddim_pipeline = DDIMPipeline(unet=model, scheduler=ddim_scheduler)
 pndm_scheduler = PNDMScheduler.from_config(model_id, subfolder="scheduler")
 pndm_pipeline = PNDMPipeline(unet=model, scheduler=pndm_scheduler)
 # run pipeline in inference (sample random noise and denoise)
-def predict(seed=42,scheduler="ddim"):
+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=100)
-        image = image["sample"]
+        image = ddim_pipeline(generator=generator, num_inference_steps=steps)["sample"]
     elif(scheduler == "ddpm"):
-        image = ddpm_pipeline(generator=generator)
-        #["sample"] doesnt work here for some reason
+        image = ddpm_pipeline(generator=generator)["sample"]
     elif(scheduler == "pndm"):
-        image = pndm_pipeline(generator=generator, num_inference_steps=11)
-        #["sample"] doesnt work here for some reason
+        image = pndm_pipeline(generator=generator, num_inference_steps=steps)["sample"]
 
     image_processed = image.cpu().permute(0, 2, 3, 1)
     if scheduler == "pndm":
@@ -49,7 +46,7 @@ random_seed = random.randint(0, 2147483647)
 gr.Interface(
     predict,
     inputs=[
-        #gr.inputs.Slider(1, 1000, label='Inference Steps', default=20, step=1),
+        gr.inputs.Slider(1, 100, label='Inference Steps', default=20, step=1),
         gr.inputs.Slider(0, 2147483647, label='Seed', default=random_seed),
         gr.inputs.Radio(["ddim", "ddpm", "pndm"], default="ddpm",label="Diffusion scheduler")
     ],