working

builder/modal-builder/src/main.py

@@ -47,6 +47,7 @@ machine_id_websocket_dict = {}
 machine_id_status = {}
 
 fly_instance_id = os.environ.get('FLY_ALLOC_ID', 'local').split('-')[0]
+civitai_api_key = os.environ.get('FLY_ALLOC_ID', 'local').split('-')[0]
 
 
 class FlyReplayMiddleware(BaseHTTPMiddleware):
@@ -272,7 +273,8 @@ async def upload_logic(body: UploadBody):
             "checkpoint_id": body.checkpoint_id,
             "volume_id": body.volume_id,
             "folder_path": upload_path,
-        }
+        },
+        "civitai_api_key": os.environ.get('CIVITAI_API_KEY')
     }
     with open(f"{folder_path}/config.py", "w") as f:
         f.write("config = " + json.dumps(config))

builder/modal-builder/src/volume-builder/app.py

@@ -25,16 +25,13 @@ vol_name_to_links = config["volume_names"]
 vol_name_to_path = config["volume_paths"]
 callback_url = config["callback_url"]
 callback_body = config["callback_body"]
+civitai_key = config["civitai_api_key"]
 
 volumes = create_volumes(vol_name_to_links, vol_name_to_path)
 image = ( 
    modal.Image.debian_slim().apt_install("wget").pip_install("requests")
 )
 
-print(vol_name_to_links)
-print(vol_name_to_path)
-print(volumes)
-
 # download config { "download_url": "", "folder_path": ""}
 timeout=5000
 @stub.function(volumes=volumes, image=image, timeout=timeout, gpu=None)
@@ -42,20 +39,23 @@ def download_model(volume_name, download_config):
     import requests
     download_url = download_config["download_url"]
     folder_path = download_config["folder_path"]
+
     volume_base_path = vol_name_to_path[volume_name]
     model_store_path = os.path.join(volume_base_path, folder_path)
+    modified_download_url = download_url + ("&" if "?" in download_url else "?") + "token=" + civitai_key
+    print('downlodaing', modified_download_url)
 
-    subprocess.run(["wget", download_url, "--content-disposition", "-P", model_store_path])
+    subprocess.run(["wget", modified_download_url , "--content-disposition", "-P", model_store_path])
     subprocess.run(["ls", "-la", volume_base_path])
     subprocess.run(["ls", "-la", model_store_path])
     volumes[volume_base_path].commit()
 
+
     status =  {"status": "success"}
     requests.post(callback_url, json={**status, **callback_body})
     print(f"finished! sending to {callback_url}")
     pprint({**status, **callback_body})
 
-
 @stub.local_entrypoint()
 def simple_download():
     import requests

builder/modal-builder/src/volume-builder/config.py

@@ -13,5 +13,6 @@ config = {
         "checkpoint_id": "",
         "volume_id": "",
         "folder_path": "images",
-    }
+    }, 
+    "civitai_api_key": "",
 }