update

app.py

@@ -131,7 +131,11 @@ def create_inference_demo(func: inference_fn) -> gr.Blocks:
                                                maximum=10.,
                                                step=1,
                                                value=10)
-
+                    ddim_steps = gr.Slider(label='Number of DDIM Sampling Steps',
+                                               minimum=10,
+                                               maximum=100,
+                                               step=1,
+                                               value=50)
                 run_button = gr.Button('Generate')
 
                 # gr.Markdown('''
@@ -146,23 +150,25 @@ def create_inference_demo(func: inference_fn) -> gr.Blocks:
         #                     inputs=None,
         #                     outputs=weight_name)
         prompt.submit(fn=func,
-                      inputs=[
-                          model_id,
-                          prompt,
-                          num_samples,
-                          guidance_scale,
-                      ],
-                      outputs=result,
-                      queue=False)
+                        inputs=[
+                            model_id,
+                            prompt,
+                            num_samples,
+                            guidance_scale,
+                            ddim_steps,
+                        ],
+                        outputs=result,
+                        queue=False)
         run_button.click(fn=func,
                         inputs=[
                             model_id,
                             prompt,
                             num_samples,
                             guidance_scale,
+                            ddim_steps,
                         ],
-                         outputs=result,
-                         queue=False)
+                        outputs=result,
+                        queue=False)
     return demo
 
 

inference.py

@@ -46,12 +46,14 @@ def inference_fn(
         prompt: str,
         num_samples: int,
         guidance_scale: float,
+        ddim_steps: int,
     ) -> PIL.Image.Image:
 
     # create inference pipeline
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    pipe = StableDiffusionPipeline.from_pretrained(os.path.join('experiments', model_id),torch_dtype=torch.float16).to(device)
-
+    if torch.cuda.is_available():
+        pipe = StableDiffusionPipeline.from_pretrained(os.path.join('experiments', model_id),torch_dtype=torch.float16).to('cuda')
+    else:
+        pipe = StableDiffusionPipeline.from_pretrained(os.path.join('experiments', model_id)).to('cpu')
     # make directory to save images
     image_root_folder = os.path.join('experiments', model_id, 'inference')
     os.makedirs(image_root_folder, exist_ok = True)
@@ -80,7 +82,7 @@ def inference_fn(
         os.makedirs(image_folder, exist_ok = True)
 
         # batch generation
-        images = pipe(prompt, num_inference_steps=50, guidance_scale=guidance_scale, num_images_per_prompt=num_samples).images
+        images = pipe(prompt, num_inference_steps=ddim_steps, guidance_scale=guidance_scale, num_images_per_prompt=num_samples).images
 
         # save generated images
         for idx, image in enumerate(images):