Update app.py

app.py

@@ -1,19 +1,20 @@
 import gradio as gr
-from Sketch2ImgPipeline import Sketch2ImgPipeline
+from pipeline_ddpm_sketch2img import DDPMSketch2ImgPipeline
 import numpy as np
+from diffusers import DDPMScheduler, DPMSolverMultistepScheduler, DDIMScheduler
+from PIL import Image
 
-pipe = Sketch2ImgPipeline.from_pretrained("IzumiSatoshi/sketch2img-FashionMNIST")
+model_path = "IzumiSatoshi/sketch2img-FashionMNIST"
+pipe = DDPMSketch2ImgPipeline.from_pretrained(model_path).to("cpu")
+pipe.scheduler = DDIMScheduler.from_pretrained(model_path, subfolder="scheduler")
 
 
-def greet(input_img):
-    sketches = np.expand_dims(input_img, (0, 1))
-    sketches[sketches < 250] = 0
-    sketches[sketches >= 250] = 255
-    print(sketches.shape)
-    samples = pipe(sketches, num_inference_step=10)
-    out = samples[0][0]
-    print(out.shape)
-    return sketches[0][0], out
+def draw(sketch):
+    sketch[sketch < 250] = 0
+    sketch[sketch >= 250] = 255
+    sketch = Image.fromarray(sketch)
+    image = pipe(sketch, num_inference_step=50)
+    return sketch, image
 
 
 inp = gr.inputs.Image(
@@ -23,5 +24,5 @@ inp = gr.inputs.Image(
     invert_colors=True,
     tool="select",
 )
-demo = gr.Interface(fn=greet, inputs=inp, outputs=["image", "image"])
+demo = gr.Interface(fn=draw, inputs=inp, outputs=["image", "image"])
 demo.launch()