some extra changes for collab compatibility

2 files changed, 12 insertions, 2 deletions

diff --git a/modules/paths.py b/modules/paths.py
index 6d11b304..e1559bc7 100644
--- a/modules/paths.py
+++ b/modules/paths.py
@@ -6,7 +6,13 @@ script_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
 sys.path.insert(0, script_path)

 

 # use current directory as SD dir if it has related files, otherwise parent dir of script as stated in guide

-sd_path = os.path.abspath('.') if os.path.exists('./ldm/models/diffusion/ddpm.py') else os.path.dirname(script_path)

+sd_path = None

+possible_sd_paths = ['.', os.path.dirname(script_path), os.path.join(script_path, 'repositories/stable-diffusion')]

+for possible_sd_path in possible_sd_paths:

+    if os.path.exists(os.path.join(possible_sd_path, 'ldm/models/diffusion/ddpm.py')):

+        sd_path = os.path.abspath(possible_sd_path)

+

+assert sd_path is not None, "Couldn't find Stable Diffusion in any of: " + possible_sd_paths

 

 # add parent directory to path; this is where Stable diffusion repo should be

 path_dirs = [

diff --git a/modules/shared.py b/modules/shared.py
index d57aba37..c86cf4a6 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -32,7 +32,11 @@ parser.add_argument("--share", action='store_true', help="use share=True for gra
 parser.add_argument("--esrgan-models-path", type=str, help="path to directory with ESRGAN models", default=os.path.join(script_path, 'ESRGAN'))

 parser.add_argument("--opt-split-attention", action='store_true', help="enable optimization that reduced vram usage by a lot for about 10% decrease in performance")

 parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")

-cmd_opts = parser.parse_args()

+

+if 'notebook_argv' in globals():

+    cmd_opts = parser.parse_args(notebook_argv)

+else:

+    cmd_opts = parser.parse_args()

 

 cpu = torch.device("cpu")

 gpu = torch.device("cuda")