diff --git a/builder/modal-builder/src/main.py b/builder/modal-builder/src/main.py index 2493a37..cc21188 100644 --- a/builder/modal-builder/src/main.py +++ b/builder/modal-builder/src/main.py @@ -312,7 +312,8 @@ async def build_logic(item: Item): config = { "name": item.name, "deploy_test": os.environ.get("DEPLOY_TEST_FLAG", "False"), - "gpu": item.gpu + "gpu": item.gpu, + "civitai_token": os.environ.get("CIVITAI_TOKEN", "") } with open(f"{folder_path}/config.py", "w") as f: f.write("config = " + json.dumps(config)) diff --git a/builder/modal-builder/src/template/app.py b/builder/modal-builder/src/template/app.py index 303bb97..1e8d5d8 100644 --- a/builder/modal-builder/src/template/app.py +++ b/builder/modal-builder/src/template/app.py @@ -36,6 +36,9 @@ if not deploy_test: dockerfile_image = ( modal.Image.debian_slim() + .env({ + "CIVITAI_TOKEN": config["civitai_token"], + }) .apt_install("git", "wget") .pip_install( "git+https://github.com/modal-labs/asgiproxy.git", "httpx", "tqdm" diff --git a/builder/modal-builder/src/template/data/install_deps.py b/builder/modal-builder/src/template/data/install_deps.py index 3ff3ca3..27ac6af 100644 --- a/builder/modal-builder/src/template/data/install_deps.py +++ b/builder/modal-builder/src/template/data/install_deps.py @@ -49,6 +49,12 @@ with open('models.json') as f: models = json.load(f) for model in models: + import os + if "civitai.com/api" in model['url'] and not "token=" in model['url']: + if "?" in model['url']: + model['url'] += "&token=" + os.environ.get('CIVITAI_TOKEN', '') + else: + model['url'] += "?token=" + os.environ.get('CIVITAI_TOKEN', '') response = requests.request("POST", f"{root_url}/model/install", json=model, headers=headers) print(response.text)