some extra changes for collab compatibility

This commit is contained in:
AUTOMATIC 2022-09-06 10:38:02 +03:00
parent 395f170597
commit c688d0984e
2 changed files with 12 additions and 2 deletions

View File

@ -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 = [

View File

@ -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")