blah

builder/modal-builder/src/main.py

@@ -8,6 +8,7 @@ from enum import Enum
 import json
 import subprocess
 import time
+from uuid import uuid4
 from contextlib import asynccontextmanager
 import asyncio
 import threading
@@ -19,6 +20,7 @@ from urllib.parse import parse_qs
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.types import ASGIApp, Scope, Receive, Send
 
+
 from concurrent.futures import ThreadPoolExecutor
 
 # executor = ThreadPoolExecutor(max_workers=5)
@@ -174,6 +176,7 @@ class Item(BaseModel):
     snapshot: Snapshot
     models: List[Model]
     callback_url: str
+    checkpoint_volume_name: str
     gpu: GPUType = Field(default=GPUType.T4)
 
     @field_validator('gpu')
@@ -223,6 +226,102 @@ async def websocket_endpoint(websocket: WebSocket, machine_id: str):
 
 #     return {"Hello": "World"}
 
+class UploadType(str, Enum):
+    checkpoint = "checkpoint"
+
+class UploadBody(BaseModel):
+    download_url: str
+    volume_name: str
+    volume_id: str
+    checkpoint_id: str
+    upload_type: UploadType
+    callback_url: str
+
+
+UPLOAD_TYPE_DIR_MAP = {
+    UploadType.checkpoint: "checkpoints"
+}
+
+@app.post("/upload-volume")
+async def upload_checkpoint(body: UploadBody):
+    global last_activity_time
+    last_activity_time = time.time()
+    logger.info(f"Extended inactivity time to {global_timeout}")
+
+    asyncio.create_task(upload_logic(body))
+
+    # check that this
+    return JSONResponse(status_code=200, content={"message": "Volume uploading", "build_machine_instance_id": fly_instance_id})
+
+async def upload_logic(body: UploadBody):
+    folder_path = f"/app/builds/{body.volume_id}"
+
+    cp_process = await asyncio.subprocess.create_subprocess_exec("cp", "-r", "/app/src/volume-builder", folder_path)
+    await cp_process.wait()
+
+    upload_path = UPLOAD_TYPE_DIR_MAP[body.upload_type]
+    config = {
+        "volume_names": {
+            body.volume_name: {"download_url": body.download_url, "folder_path": upload_path}
+        }, 
+        "volume_paths": {
+            body.volume_name: f'/volumes/{uuid4()}'
+        }, 
+        "callback_url": body.callback_url,
+        "callback_body": {
+            "checkpoint_id": body.checkpoint_id,
+            "volume_id": body.volume_id,
+            "folder_path": upload_path,
+        }
+    }
+    with open(f"{folder_path}/config.py", "w") as f:
+        f.write("config = " + json.dumps(config))
+
+    process = await asyncio.subprocess.create_subprocess_shell(
+        f"modal run app.py",
+        # stdout=asyncio.subprocess.PIPE,
+        # stderr=asyncio.subprocess.PIPE,
+        cwd=folder_path,
+        env={**os.environ, "COLUMNS": "10000"}
+    )
+    
+    # error_logs = []
+    # async def read_stream(stream):
+    #     while True:
+    #         line = await stream.readline()
+    #         if line:
+    #             l = line.decode('utf-8').strip()
+    #             error_logs.append(l)
+    #             logger.error(l)
+    #             error_logs.append({
+    #                 "logs": l,
+    #                 "timestamp": time.time()
+    #             })
+    #         else:
+    #             break
+
+    # stderr_read_task = asyncio.create_task(read_stream(process.stderr))
+    #
+    # await asyncio.wait([stderr_read_task])
+    # await process.wait()
+
+    # if process.returncode != 0:
+    #     error_logs.append({"logs": "Unable to upload volume.", "timestamp": time.time()})
+    #     # Error handling: send POST request to callback URL with error details
+    #     requests.post(body.callback_url, json={
+    #         "volume_id": body.volume_id, 
+    #         "checkpoint_id": body.checkpoint_id,
+    #         "folder_path": upload_path,
+    #         "error_logs": json.dumps(error_logs),
+    #         "status": "failed"
+    #     })
+    #
+    # requests.post(body.callback_url, json={
+    #     "checkpoint_id": body.checkpoint_id,
+    #     "volume_id": body.volume_id,
+    #     "folder_path": upload_path,
+    #     "status": "success"
+    # })
 
 @app.post("/create")
 async def create_machine(item: Item):
@@ -312,7 +411,9 @@ 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,
+        "public_checkpoint_volume": "model-store",
+        "private_checkpoint_volume": item.checkpoint_volume_name
     }
     with open(f"{folder_path}/config.py", "w") as f:
         f.write("config = " + json.dumps(config))

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

@@ -7,6 +7,7 @@ import urllib.parse
 from pydantic import BaseModel
 from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse
+from volume_setup import volumes
 
 # deploy_test = False
 
@@ -27,6 +28,7 @@ deploy_test = config["deploy_test"] == "True"
 web_app = FastAPI()
 print(config)
 print("deploy_test ", deploy_test)
+print('volumes', volumes)
 stub = Stub(name=config["name"])
 # print(stub.app_id)
 
@@ -56,7 +58,7 @@ if not deploy_test:
         #     # Install comfy deploy
         #     "cd /comfyui/custom_nodes && git clone https://github.com/BennyKok/comfyui-deploy.git",
         # )
-        # .copy_local_file(f"{current_directory}/data/extra_model_paths.yaml", "/comfyui")
+        .copy_local_file(f"{current_directory}/data/extra_model_paths.yaml", "/comfyui")
 
         .copy_local_file(f"{current_directory}/data/start.sh", "/start.sh")
         .run_commands("chmod +x /start.sh")
@@ -154,7 +156,7 @@ image = Image.debian_slim()
 target_image = image if deploy_test else dockerfile_image
 
 
-@stub.function(image=target_image, gpu=config["gpu"])
+@stub.function(image=target_image, gpu=config["gpu"], volumes=volumes)
 def run(input: Input):
     import subprocess
     import time
@@ -235,7 +237,7 @@ async def bar(request_input: RequestInput):
     # pass
 
 
-@stub.function(image=image)
+@stub.function(image=image, volumes=volumes)
 @asgi_app()
 def comfyui_api():
     return web_app
@@ -284,6 +286,7 @@ def spawn_comfyui_in_background():
     # Restrict to 1 container because we want to our ComfyUI session state
     # to be on a single container.
     concurrency_limit=1,
+    volumes=volumes,
     timeout=10 * 60,
 )
 @asgi_app()

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

@@ -1 +1,7 @@
-config = {"name": "my-app", "deploy_test": "True", "gpu": "T4"}
+config = {
+    "name": "my-app",
+    "deploy_test": "True",
+    "gpu": "T4", 
+    "public_checkpoint_volume": "model-store",
+    "private_checkpoint_volume": "private-model-store"
+}

builder/modal-builder/src/template/data/extra_model_paths.yaml

@@ -1,11 +1,15 @@
-comfyui:
-    base_path: /runpod-volume/ComfyUI/
-    checkpoints: models/checkpoints/
-    clip: models/clip/
-    clip_vision: models/clip_vision/
-    configs: models/configs/
-    controlnet: models/controlnet/
-    embeddings: models/embeddings/
-    loras: models/loras/
-    upscale_models: models/upscale_models/
-    vae: models/vae/
+public:
+  base_path: /public_models/
+  checkpoints: checkpoints
+  clip: clip
+  clip_vision: clip_vision
+  configs: configs
+  controlnet: controlnet
+  embeddings: embeddings
+  loras: loras
+  upscale_models: upscale_models
+  vae: vae
+
+private:
+  base_path: /private_models/
+  checkpoints: checkpoints

builder/modal-builder/src/template/data/insert_models.py

@@ -0,0 +1,101 @@
+"""
+This is a standalone script to download models into a modal Volume using civitai
+
+Example Usage
+`modal run insert_models::insert_model --civitai-url https://civitai.com/models/36520/ghostmix`
+This inserts an individual model from a civitai url 
+
+`modal run insert_models::insert_models_civitai_api` 
+This inserts a bunch of models based on the models retrieved by civitai
+
+civitai's API reference https://github.com/civitai/civitai/wiki/REST-API-Reference
+"""
+import modal
+import subprocess
+import requests
+import json
+
+stub = modal.Stub()
+
+# NOTE: volume name can be variable
+volume = modal.Volume.persisted("rah")
+model_store_path = "/vol/models"
+MODEL_ROUTE = "models"
+image = (
+    modal.Image.debian_slim().apt_install("wget").pip_install("requests")
+)
+
+@stub.function(volumes={model_store_path: volume}, image=image, timeout=50000, gpu=None)
+def download_model(download_url):
+    print(download_url)
+    subprocess.run(["wget", download_url, "--content-disposition", "-P", model_store_path])
+    subprocess.run(["ls", "-la", model_store_path])
+    volume.commit()
+
+# file is raw output from Civitai API https://github.com/civitai/civitai/wiki/REST-API-Reference
+
+@stub.function()
+def get_civitai_models(model_type: str, sort: str = "Highest Rated", page: int = 1):
+    """Fetch models from CivitAI API based on type."""
+    try:
+        response = requests.get(f"https://civitai.com/api/v1/models", params={"types": model_type, "page": page, "sort": sort})
+        response.raise_for_status()
+        return response.json()
+    except requests.RequestException as e:
+        print(f"Error fetching models: {e}")
+        return None
+
+
+@stub.function()
+def get_civitai_model_url(civitai_url: str):
+    # Validate the URL
+
+    if civitai_url.startswith("https://civitai.com/api/"):
+        api_url = civitai_url
+    elif civitai_url.startswith("https://civitai.com/models/"):  
+        try:
+            model_id = civitai_url.split("/")[4]
+            int(model_id) 
+        except (IndexError, ValueError):
+            return None 
+        api_url = f"https://civitai.com/api/v1/models/{model_id}"
+    else:
+        return "Error: URL must be from civitai.com and contain /models/"
+
+    response = requests.get(api_url)
+    # Check for successful response
+    if response.status_code != 200:
+        return f"Error: Unable to fetch data from {api_url}"
+    # Return the response data
+    return response.json()
+
+
+
+@stub.local_entrypoint()
+def insert_models_civitai_api(type: str = "Checkpoint", sort = "Highest Rated", page: int = 1):
+    civitai_models = get_civitai_models.local(type, sort, page)
+    if civitai_models:
+        for _ in download_model.map(map(lambda model: model['modelVersions'][0]['downloadUrl'], civitai_models['items'])):
+            pass
+    else:
+        print("Failed to retrieve models.")
+
+@stub.local_entrypoint()
+def insert_model(civitai_url: str):
+    if civitai_url.startswith("'https://civitai.com/api/download/models/"):
+        download_url = civitai_url
+    else:
+        civitai_model = get_civitai_model_url.local(civitai_url)
+        if civitai_model:
+            download_url = civitai_model['modelVersions'][0]['downloadUrl']
+        else:
+            return "invalid URL"
+
+    download_model.remote(download_url)
+
+@stub.local_entrypoint()
+def simple_download():
+    download_urls = ['https://civitai.com/api/download/models/119057', 'https://civitai.com/api/download/models/130090', 'https://civitai.com/api/download/models/31859', 'https://civitai.com/api/download/models/128713', 'https://civitai.com/api/download/models/179657', 'https://civitai.com/api/download/models/143906', 'https://civitai.com/api/download/models/9208', 'https://civitai.com/api/download/models/136078', 'https://civitai.com/api/download/models/134065', 'https://civitai.com/api/download/models/288775', 'https://civitai.com/api/download/models/95263', 'https://civitai.com/api/download/models/288982', 'https://civitai.com/api/download/models/87153', 'https://civitai.com/api/download/models/10638', 'https://civitai.com/api/download/models/263809', 'https://civitai.com/api/download/models/130072', 'https://civitai.com/api/download/models/117019', 'https://civitai.com/api/download/models/95256', 'https://civitai.com/api/download/models/197181', 'https://civitai.com/api/download/models/256915', 'https://civitai.com/api/download/models/118945', 'https://civitai.com/api/download/models/125843', 'https://civitai.com/api/download/models/179015', 'https://civitai.com/api/download/models/245598', 'https://civitai.com/api/download/models/223670', 'https://civitai.com/api/download/models/90072', 'https://civitai.com/api/download/models/290817', 'https://civitai.com/api/download/models/154097', 'https://civitai.com/api/download/models/143497', 'https://civitai.com/api/download/models/5637']
+
+    for _ in download_model.map(download_urls):
+        pass

builder/modal-builder/src/template/data/install_deps.py

@@ -45,12 +45,12 @@ for package in packages:
     response = requests.request("POST", f"{root_url}/customnode/install", json=package, headers=headers)
     print(response.text)
 
-with open('models.json') as f:
-    models = json.load(f)
+# with open('models.json') as f:
+#     models = json.load(f)
 
-for model in models:
-    response = requests.request("POST", f"{root_url}/model/install", json=model, headers=headers)
-    print(response.text)
+# for model in models:
+#     response = requests.request("POST", f"{root_url}/model/install", json=model, headers=headers)
+#     print(response.text)
 
 # Close the server
 server_process.terminate()

builder/modal-builder/src/template/volume_setup.py

@@ -0,0 +1,9 @@
+import modal
+from config import config
+
+public_model_volume = modal.Volume.persisted(config["public_checkpoint_volume"])
+private_volume = modal.Volume.persisted(config["private_checkpoint_volume"])
+
+PUBLIC_BASEMODEL_DIR = "/public_models"
+PRIVATE_BASEMODEL_DIR = "/private_models"
+volumes = {PUBLIC_BASEMODEL_DIR: public_model_volume, PRIVATE_BASEMODEL_DIR: private_volume}

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

@@ -0,0 +1,69 @@
+import modal
+from config import config
+import os
+import subprocess
+
+stub = modal.Stub()
+
+# Volume names may only contain alphanumeric characters, dashes, periods, and underscores, and must be less than 64 characters in length.
+def is_valid_name(name: str) -> bool:
+    allowed_characters = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-._")
+    return 0 < len(name) <= 64 and all(char in allowed_characters for char in name)
+
+def create_volumes(volume_names, paths):
+    path_to_vol = {}
+    for volume_name in volume_names.keys():
+        if not is_valid_name(volume_name):
+            pass
+        modal_volume = modal.Volume.persisted(volume_name)
+        path_to_vol[paths[volume_name]] = modal_volume
+ 
+    return path_to_vol
+
+vol_name_to_links = config["volume_names"]
+vol_name_to_path = config["volume_paths"]
+callback_url = config["callback_url"]
+callback_body = config["callback_body"]
+
+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)
+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)
+
+    subprocess.run(["wget", 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})
+
+
+@stub.local_entrypoint()
+def simple_download():
+    import requests
+    print(vol_name_to_links)
+    print([(vol_name, link) for vol_name,link in vol_name_to_links.items()])
+    try:
+        list(download_model.starmap([(vol_name, link) for vol_name,link in vol_name_to_links.items()]))
+    except modal.exception.FunctionTimeoutError as e:
+        status =  {"status": "failed", "error_logs": f"{str(e)}", "timeout": timeout}
+        requests.post(callback_url, json={**status, **callback_body})
+    except Exception as e:
+        status =  {"status": "failed", "error_logs": str(e)}
+        requests.post(callback_url, json={**status, **callback_body})
+        

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

@@ -0,0 +1,17 @@
+config = {
+    "volume_names": {
+        "test": {
+            "download_url": "https://pub-6230db03dc3a4861a9c3e55145ceda44.r2.dev/openpose-pose (1).png",
+            "folder_path": "images"
+        }
+    }, 
+    "volume_paths": {
+        "test": "/volumes/something"
+    },
+    "callback_url": "",
+    "callback_body": {
+        "checkpoint_id": "",
+        "volume_id": "",
+        "folder_path": "images",
+    }
+}