blah

2024-01-24 01:58:00 -08:00
parent 6437de4def
commit b1e9bcc4e6
25 changed files with 2561 additions and 30 deletions
@@ -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()
@@ -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"
+}
@@ -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
@@ -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
@@ -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()
@@ -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}