update

app.py

@@ -101,7 +101,7 @@ def reload_custom_diffusion_weight_list() -> dict:
     return gr.update(choices=find_weight_files())
 
 
-def create_inference_demo(func: inference_fn, device) -> gr.Blocks:
+def create_inference_demo(func: inference_fn) -> gr.Blocks:
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
@@ -151,7 +151,6 @@ def create_inference_demo(func: inference_fn, device) -> gr.Blocks:
                           prompt,
                           num_samples,
                           guidance_scale,
-                          device
                       ],
                       outputs=result,
                       queue=False)
@@ -161,7 +160,6 @@ def create_inference_demo(func: inference_fn, device) -> gr.Blocks:
                             prompt,
                             num_samples,
                             guidance_scale,
-                            device
                         ],
                          outputs=result,
                          queue=False)
@@ -184,7 +182,7 @@ with gr.Blocks(css='style.css') as demo:
 
     with gr.Tabs():
         with gr.TabItem('Test'):
-            create_inference_demo(inference_fn, device=args.device)
+            create_inference_demo(inference_fn)
 
 demo.launch(
     enable_queue=args.enable_queue,

inference.py

@@ -46,10 +46,10 @@ def inference_fn(
         prompt,
         num_samples,
         guidance_scale,
-        device
     ):
 
     # create inference pipeline
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
     pipe = StableDiffusionPipeline.from_pretrained(model_id,torch_dtype=torch.float16).to(device)
 
     # make directory to save images