edit device; edit dropdown

app.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python
 """Demo app for https://github.com/ziqihuangg/ReVersion.
 The code in this repo is partly adapted from the following repository:
-https://huggingface.co/spaces/hysts/LoRA-SD-training
+https://huggingface.co/spaces/nupurkmr9/custom-diffusion
 
 S-Lab License 1.0
 
-Copyright 2022 S-Lab
+Copyright 2023 S-Lab
 Redistribution and use for non-commercial purpose in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
@@ -101,13 +101,13 @@ def reload_custom_diffusion_weight_list() -> dict:
     return gr.update(choices=find_weight_files())
 
 
-def create_inference_demo(func: inference_fn) -> gr.Blocks:
+def create_inference_demo(func: inference_fn, device=args.device) -> gr.Blocks:
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 model_id = gr.Dropdown(
-                    choices=['experiments/painted_on'],
-                    value='experiments/painted_on',
+                    choices=['painted_on'],
+                    value='painted_on',
                     label='Relation',
                     visible=True)
                 # reload_button = gr.Button('Reload Weight List')
@@ -151,7 +151,8 @@ def create_inference_demo(func: inference_fn) -> gr.Blocks:
                           prompt,
                           placeholder_string,
                           num_samples,
-                          guidance_scale
+                          guidance_scale,
+                          device
                       ],
                       outputs=result,
                       queue=False)
@@ -161,7 +162,8 @@ def create_inference_demo(func: inference_fn) -> gr.Blocks:
                             prompt,
                             placeholder_string,
                             num_samples,
-                            guidance_scale
+                            guidance_scale,
+                            device
                         ],
                          outputs=result,
                          queue=False)
@@ -184,7 +186,7 @@ with gr.Blocks(css='style.css') as demo:
 
     with gr.Tabs():
         with gr.TabItem('Test'):
-            create_inference_demo(inference_fn)
+            create_inference_demo(inference_fn, args.device)
 
 demo.launch(
     enable_queue=args.enable_queue,

inference.py

@@ -46,14 +46,15 @@ def inference_fn(
         prompt,
         placeholder_string,
         num_samples,
-        guidance_scale
+        guidance_scale,
+        device
     ):
 
     # create inference pipeline
-    pipe = StableDiffusionPipeline.from_pretrained(model_id,torch_dtype=torch.float16).to("cuda")
+    pipe = StableDiffusionPipeline.from_pretrained(model_id,torch_dtype=torch.float16).to(device)
 
     # make directory to save images
-    image_root_folder = os.path.join(model_id, 'inference')
+    image_root_folder = os.path.join('experiments', model_id, 'inference')
     os.makedirs(image_root_folder, exist_ok = True)
 
     # if prompt is None and args.template_name is None: