chore: update default git url

fix: async issues with updating python to 3.11.1
fix: revert app back with async
2024-02-02 00:55:38 +08:00 · 2024-02-02 00:31:04 +08:00 · 2024-02-02 00:11:40 +08:00 · 2024-02-01 23:38:35 +08:00 · 2024-02-01 22:11:17 +08:00 · 2024-02-01 19:42:46 +08:00
79 changed files with 7966 additions and 1419 deletions
@@ -13,4 +13,4 @@ RUN mkdir builds
 # CMD ["uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "80", "--lifespan", "on"]
 CMD ["python", "src/main.py"]
 # If running behind a proxy like Nginx or Traefik add --proxy-headers
-# CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "80", "--proxy-headers"]
+# CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "80", "--proxy-headers"]
@@ -56,5 +56,5 @@ fly launch
 ```
 if not, run this instead
 ```
-fly deploy
+fly deploy -c "toml file"
 ```
@@ -19,7 +19,7 @@ import requests
 from urllib.parse import parse_qs
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.types import ASGIApp, Scope, Receive, Send
-
+import modal

 from concurrent.futures import ThreadPoolExecutor

@@ -180,6 +180,8 @@ class Item(BaseModel):
    models: List[Model]
    callback_url: str
    model_volume_name: str
+    run_timeout: Optional[int] = Field(default=60 * 5)
+    idle_timeout: Optional[int] = Field(default=60)
    gpu: GPUType = Field(default=GPUType.T4)

    @field_validator('gpu')
@@ -234,6 +236,14 @@ class UploadType(str, Enum):
    checkpoint = "checkpoint"
    lora = "lora"
    embedding = "embedding"
+    clip = "clip"
+    clip_vision = "clip_vision"
+    configs = "configs"
+    controlnet = "controlnet"
+    upscale_models = "upscale_models"
+    vae = "vae"
+    ipadapter = "ipadapter"
+    other = "other"

 class UploadBody(BaseModel):
    download_url: str
@@ -249,8 +259,46 @@ UPLOAD_TYPE_DIR_MAP = {
    UploadType.checkpoint: "checkpoints",
    UploadType.lora: "loras",
    UploadType.embedding: "embeddings",
+    UploadType.clip: "clip",
+    UploadType.clip_vision: "clip_vision",
+    UploadType.configs: "configs",
+    UploadType.controlnet: "controlnet",
+    UploadType.upscale_models: "upscale_models",
+    UploadType.vae: "vae",
+    UploadType.ipadapter: "ipadapter",
+    UploadType.other: "",
 }

+class DeleteBody(BaseModel):
+    volume_name: str
+    path: str
+    file_name: str
+
+
+@app.post("/delete-volume-model")
+async def delete_model(body: DeleteBody):
+    global last_activity_time
+    last_activity_time = time.time()
+    logger.info(f"Extended inactivity time to {global_timeout}")
+    
+    full_path = f"{body.path.rstrip('/')}/{body.file_name}"
+
+    rm_process = await asyncio.subprocess.create_subprocess_exec("modal", "volume", "rm", body.volume_name, full_path,
+                                                                 stdout=asyncio.subprocess.PIPE,
+                                                                 stderr=asyncio.subprocess.PIPE,)
+    await rm_process.wait()
+
+    logger.info(f"Successfully deleted: {full_path} from volume: {body.volume_name}")
+    stdout, stderr = await rm_process.communicate()
+    if stdout:
+        logger.info(f"cp_process stdout: {stdout.decode()}")
+    if stderr:
+        logger.info(f"cp_process stderr: {stderr.decode()}")
+
+    if rm_process.returncode == 0:
+        return JSONResponse(status_code=200, content={"status":f"Successfully deleted {full_path} from volume {body.volume_name}"})
+    else:
+        return JSONResponse(status_code=500, content={"status": "error", "error": stderr.decode()})

@app.post("/upload-volume")
 async def upload_model(body: UploadBody):
@@ -265,12 +313,16 @@ async def upload_model(body: UploadBody):


 async def upload_logic(body: UploadBody):
-    folder_path = f"/app/builds/{body.volume_id}"
+    folder_path = f"/app/builds/{body.volume_id}-{uuid4()}"

-    cp_process = await asyncio.subprocess.create_subprocess_exec("cp", "-r", "/app/src/volume-builder", folder_path)
+    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]
+    if upload_path == "":
+        # TODO: deal with custom paths
+        pass
+
    config = {
        "volume_names": {
            body.volume_name: {"download_url": body.download_url, "folder_path": upload_path}
@@ -284,16 +336,22 @@ async def upload_logic(body: UploadBody):
            "volume_id": body.volume_id,
            "folder_path": upload_path,
        },
-        "civitai_api_key": os.environ.get('CIVITAI_API_KEY')
+        "civitai_api_key": os.environ.get('CIVITAI_API_KEY'),
+        "app_name": f"vol_name_{uuid4()}"
    }
    with open(f"{folder_path}/config.py", "w") as f:
        f.write("config = " + json.dumps(config))

-    await asyncio.subprocess.create_subprocess_shell(
-        f"modal run app.py",
+    process = await asyncio.subprocess.create_subprocess_shell(
+        f"python runner.py",
        cwd=folder_path,
        env={**os.environ, "COLUMNS": "10000"}
    )
+    await process.wait()
+
+    # import modal
+    # modal.deploy_stub(stub)
+    # stub["download_model"].web_url

@app.post("/create")
 async def create_machine(item: Item):
@@ -391,7 +449,9 @@ async def build_logic(item: Item):
        "gpu": item.gpu,
        "public_model_volume": public_model_volume_name,
        "private_model_volume": item.model_volume_name,
-        "pip": list(pip_modules)
+        "pip": list(pip_modules),
+        "run_timeout": item.run_timeout,
+        "idle_timeout": item.idle_timeout,
    }
    with open(f"{folder_path}/config.py", "w") as f:
        f.write("config = " + json.dumps(config))
@@ -9,6 +9,8 @@ from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import HTMLResponse
 from volume_setup import volumes
 from datetime import datetime
+import aiohttp
+from aiohttp import TCPConnector
 # deploy_test = False

 import os
@@ -36,7 +38,9 @@ if not deploy_test:
    # dockerfile_image = Image.from_dockerfile(f"{current_directory}/Dockerfile", context_mount=Mount.from_local_dir(f"{current_directory}/data", remote_path="/data"))

    dockerfile_image = (
-        modal.Image.debian_slim()
+        modal.Image.debian_slim(
+            python_version="3.11.1"
+        )
        .apt_install("git", "wget")
        .pip_install(
            "git+https://github.com/modal-labs/asgiproxy.git", "httpx", "tqdm"
@@ -49,12 +53,12 @@ if not deploy_test:

            # Install comfyui manager
            "cd /comfyui/custom_nodes && git clone https://github.com/ltdrdata/ComfyUI-Manager.git",
-            "cd /comfyui/custom_nodes/ComfyUI-Manager && git reset --hard 9c86f62b912f4625fe2b929c7fc61deb9d16f6d3",
+            "cd /comfyui/custom_nodes/ComfyUI-Manager && git reset --hard 8dd801435bb75aa1d24b7e382bac070a4c18bc51",
            "cd /comfyui/custom_nodes/ComfyUI-Manager && pip install -r requirements.txt",
            "cd /comfyui/custom_nodes/ComfyUI-Manager && mkdir startup-scripts",
        )
-        .run_commands(f"cat /comfyui/server.py")
-        .run_commands(f"ls /comfyui/app")
+        # .run_commands(f"cat /comfyui/server.py")
+        # .run_commands(f"ls /comfyui/app")
        # .run_commands(
        #     # Install comfy deploy
        #     "cd /comfyui/custom_nodes && git clone https://github.com/BennyKok/comfyui-deploy.git",
@@ -82,7 +86,7 @@ if not deploy_test:
 # Time to wait between API check attempts in milliseconds
 COMFY_API_AVAILABLE_INTERVAL_MS = 50
 # Maximum number of API check attempts
-COMFY_API_AVAILABLE_MAX_RETRIES = 500
+COMFY_API_AVAILABLE_MAX_RETRIES = 1000
 # Time to wait between poll attempts in milliseconds
 COMFY_POLLING_INTERVAL_MS = 250
 # Maximum number of poll attempts
@@ -91,48 +95,34 @@ COMFY_POLLING_MAX_RETRIES = 1000
 COMFY_HOST = "127.0.0.1:8188"


-def check_server(url, retries=50, delay=500):
-    import requests
-    import time
-    """
-    Check if a server is reachable via HTTP GET request
-
-    Args:
-    - url (str): The URL to check
-    - retries (int, optional): The number of times to attempt connecting to the server. Default is 50
-    - delay (int, optional): The time in milliseconds to wait between retries. Default is 500
-
-    Returns:
-    bool: True if the server is reachable within the given number of retries, otherwise False
-    """
-
-    for i in range(retries):
+async def check_server(url, retries=50, delay=500):
+    import aiohttp
+    # for i in range(retries):
+    while True:
        try:
-            response = requests.get(url)
-
-            # If the response status code is 200, the server is up and running
-            if response.status_code == 200:
-                print(f"runpod-worker-comfy - API is reachable")
-                return True
-        except requests.RequestException as e:
+            async with aiohttp.ClientSession() as session:
+                async with session.get(url) as response:
+                    # If the response status code is 200, the server is up and running
+                    if response.status == 200:
+                        print(f"comfy-modal - API is reachable")
+                        return True
+        except Exception as e:
            # If an exception occurs, the server may not be ready
            pass

-        # print(f"runpod-worker-comfy - trying")
-
        # Wait for the specified delay before retrying
-        time.sleep(delay / 1000)
+        await asyncio.sleep(delay / 1000)

    print(
-        f"runpod-worker-comfy - Failed to connect to server at {url} after {retries} attempts."
+        f"comfy-modal - Failed to connect to server at {url} after {retries} attempts."
    )
    return False


-def check_status(prompt_id):
-    req = urllib.request.Request(
-        f"http://{COMFY_HOST}/comfyui-deploy/check-status?prompt_id={prompt_id}")
-    return json.loads(urllib.request.urlopen(req).read())
+async def check_status(prompt_id):
+    async with aiohttp.ClientSession() as session:
+        async with session.get(f"http://{COMFY_HOST}/comfyui-deploy/check-status?prompt_id={prompt_id}") as response:
+            return await response.json()


 class Input(BaseModel):
@@ -142,12 +132,12 @@ class Input(BaseModel):
    file_upload_endpoint: str


-def queue_workflow_comfy_deploy(data: Input):
+async def queue_workflow_comfy_deploy(data: Input):
    data_str = data.json()
    data_bytes = data_str.encode('utf-8')
-    req = urllib.request.Request(
-        f"http://{COMFY_HOST}/comfyui-deploy/run", data=data_bytes)
-    return json.loads(urllib.request.urlopen(req).read())
+    async with aiohttp.ClientSession() as session:
+        async with session.post(f"http://{COMFY_HOST}/comfyui-deploy/run", data=data_bytes) as response:
+            return await response.json()


 class RequestInput(BaseModel):
@@ -158,85 +148,226 @@ image = Image.debian_slim()

 target_image = image if deploy_test else dockerfile_image

-@stub.cls(image=target_image, gpu=config["gpu"] ,volumes=volumes, timeout=60 * 10, container_idle_timeout=60 * 5)
+run_timeout = config["run_timeout"]
+idle_timeout = config["idle_timeout"]
+
+import asyncio
+
+@stub.cls(
+    image=target_image, 
+    gpu=config["gpu"] ,
+    volumes=volumes, 
+    timeout=(config["run_timeout"] + 20), 
+    container_idle_timeout=config["idle_timeout"],
+    allow_concurrent_inputs=1,
+)
 class ComfyDeployRunner:

+    machine_logs = []
+
+    async def read_stream(self, stream, isStderr):
+        import time
+        while True:
+            try:
+                line = await stream.readline()
+                if line:
+                    l = line.decode('utf-8').strip()
+
+                    if l == "":
+                        continue
+
+                    if not isStderr:
+                        print(l, flush=True)
+                        self.machine_logs.append({
+                            "logs": l,
+                            "timestamp": time.time()
+                        })
+
+                    else:
+                        # is error
+                        # logger.error(l)
+                        print(l, flush=True)
+                        self.machine_logs.append({
+                            "logs": l,
+                            "timestamp": time.time()
+                        })
+                else:
+                    break
+            except asyncio.CancelledError:
+                # Handle the cancellation here if needed
+                break  # Break out of the loop on cancellation
+
    @enter()
-    def setup(self):
+    async def setup(self):
        import subprocess
        import time
+
        # Make sure that the ComfyUI API is available
        print(f"comfy-modal - check server")

-        command = ["python", "main.py",
-                "--disable-auto-launch", "--disable-metadata"]
-
-        self.server_process = subprocess.Popen(command, cwd="/comfyui")
-
-        check_server(
-            f"http://{COMFY_HOST}",
-            COMFY_API_AVAILABLE_MAX_RETRIES,
-            COMFY_API_AVAILABLE_INTERVAL_MS,
+        self.server_process = await asyncio.subprocess.create_subprocess_shell(
+            f"python main.py --disable-auto-launch --disable-metadata",
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd="/comfyui",
+            # env={**os.environ, "COLUMNS": "10000"}
        )

    @exit()
-    def cleanup(self, exc_type, exc_value, traceback):
-        self.server_process.terminate()
+    async def cleanup(self, exc_type, exc_value, traceback):
+        print(f"comfy-modal - cleanup", exc_type, exc_value, traceback)
+        # Get the current event loop
+        loop = asyncio.get_event_loop()
+
+        # Check if the event loop is closed
+        if loop.is_closed():
+            print("The event loop is closed.")
+        else:
+            try:
+                self.server_process.terminate()
+                await self.server_process.wait()
+            except Exception as e:
+                print("Issues when cleaning up", e)
+            print("The event loop is open.")

    @method()
-    def run(self, input: Input):
-        data = json.dumps({
-            "run_id": input.prompt_id,
-            "status": "started",
-            "time": datetime.now().isoformat()
-        }).encode('utf-8')
-        req = urllib.request.Request(input.status_endpoint, data=data, method='POST')
-        urllib.request.urlopen(req)
-
-        job_input = input
+    async def run(self, input: Input):
+        import signal
+        import time
+        import aiohttp

+        stdout_task = asyncio.create_task(
+            self.read_stream(self.server_process.stdout, False))
+        stderr_task = asyncio.create_task(
+            self.read_stream(self.server_process.stderr, True))
+        
        try:
-            queued_workflow = queue_workflow_comfy_deploy(job_input)  # queue_workflow(workflow)
-            prompt_id = queued_workflow["prompt_id"]
-            print(f"comfy-modal - queued workflow with ID {prompt_id}")
-        except Exception as e:
-            import traceback
-            print(traceback.format_exc())
-            return {"error": f"Error queuing workflow: {str(e)}"}
+            class TimeoutError(Exception):
+                pass

-        # Poll for completion
-        print(f"comfy-modal - wait until image generation is complete")
-        retries = 0
-        status = ""
-        try:
-            print("getting request")
-            while retries < COMFY_POLLING_MAX_RETRIES:
-                status_result = check_status(prompt_id=prompt_id)
-                # history = get_history(prompt_id)
+            def timeout_handler(signum, frame):
+                data = json.dumps({
+                    "run_id": input.prompt_id,
+                    "status": "timeout",
+                    "time": datetime.now().isoformat()
+                }).encode('utf-8')
+                req = urllib.request.Request(input.status_endpoint, data=data, method='POST')
+                urllib.request.urlopen(req)
+                raise TimeoutError("Operation timed out")
+            
+            signal.signal(signal.SIGALRM, timeout_handler)

-                # Exit the loop if we have found the history
-                # if prompt_id in history and history[prompt_id].get("outputs"):
-                #     break
+            try:
+                signal.alarm(run_timeout)

-                # Exit the loop if we have found the status both success or failed
-                if 'status' in status_result and (status_result['status'] == 'success' or status_result['status'] == 'failed'):
-                    status = status_result['status']
-                    print(status)
-                    break
-                else:
-                    # Wait before trying again
-                    time.sleep(COMFY_POLLING_INTERVAL_MS / 1000)
-                    retries += 1
-            else:
-                return {"error": "Max retries reached while waiting for image generation"}
-        except Exception as e:
-            return {"error": f"Error waiting for image generation: {str(e)}"}
+                ok = await check_server(
+                    f"http://{COMFY_HOST}",
+                    COMFY_API_AVAILABLE_MAX_RETRIES,
+                    COMFY_API_AVAILABLE_INTERVAL_MS,
+                )

-        print(f"comfy-modal - Finished, turning off")
+                if not ok:
+                    raise Exception("ComfyUI API is not available")
+                # Set an alarm for some seconds in the future

-        result = {"status": status}
+                data = json.dumps({
+                    "run_id": input.prompt_id,
+                    "status": "started",
+                    "time": datetime.now().isoformat()
+                }).encode('utf-8')
+                async with aiohttp.ClientSession() as session:
+                    async with session.post(input.status_endpoint, data=data) as response:
+                        pass

+                job_input = input
+
+                try:
+                    queued_workflow = await queue_workflow_comfy_deploy(job_input)  # queue_workflow(workflow)
+                    prompt_id = queued_workflow["prompt_id"]
+                    print(f"comfy-modal - queued workflow with ID {prompt_id}")
+                except Exception as e:
+                    import traceback
+                    print(traceback.format_exc())
+                    return {"error": f"Error queuing workflow: {str(e)}"}
+
+                # Poll for completion
+                print(f"comfy-modal - wait until image generation is complete")
+                retries = 0
+                status = ""
+                try:
+                    print("getting request")
+                    # while retries < COMFY_POLLING_MAX_RETRIES:
+                    while True:
+                        status_result = await check_status(prompt_id=prompt_id)
+                        if 'status' in status_result and (status_result['status'] == 'success' or status_result['status'] == 'failed'):
+                            status = status_result['status']
+                            print(status)
+                            break
+                        else:
+                            # Wait before trying again
+                            await asyncio.sleep(COMFY_POLLING_INTERVAL_MS / 1000)
+                            retries += 1
+                    else:
+                        return {"error": "Max retries reached while waiting for image generation"}
+                except Exception as e:
+                    return {"error": f"Error waiting for image generation: {str(e)}"}
+
+                print(f"comfy-modal - Finished, turning off")
+
+                result = {"status": status}
+
+            except TimeoutError:
+                print("Operation timed out")
+                return {"status": "failed"}
+            except Exception as e:
+                print(f"Unexpected error occurred: {str(e)}")
+                data = json.dumps({
+                    "run_id": input.prompt_id,
+                    "status": "failed",
+                    "time": datetime.now().isoformat()
+                }).encode('utf-8')
+                async with aiohttp.ClientSession() as session:
+                    async with session.post(input.status_endpoint, data=data) as response:
+                        print("response", response)
+                self.machine_logs.append({
+                    "logs": str(e),
+                    "timestamp": time.time()
+                })
+            finally:
+                signal.alarm(0)
+                
+            print("uploading log_data")
+            data = json.dumps({
+                "run_id": input.prompt_id,
+                "time": datetime.now().isoformat(),
+                "log_data": self.machine_logs
+            }).encode('utf-8')
+            print("my logs", len(self.machine_logs))
+            # Clear logs
+            timeout = aiohttp.ClientTimeout(total=60)  # 60 seconds total timeout
+            # Use HTTP/1.1 explicitly and increase the connection pool size
+            connector = TCPConnector(limit=100, force_close=True, enable_cleanup_closed=True)
+            
+            async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
+                try:
+                    async with session.post(input.status_endpoint, data=data) as response:
+                        print("response", response)
+                        # Process your response here
+                except asyncio.TimeoutError:
+                    print("Request timed out")
+                except Exception as e:
+                    print(f"An error occurred: {e}")
+            print("uploaded log_data")
+            # print(data)
+            self.machine_logs = []
+        finally:
+            stdout_task.cancel()
+            stderr_task.cancel()
+            await stdout_task
+            await stderr_task
+        
        return result
+        
    

@web_app.post("/run")
@@ -252,10 +383,12 @@ async def post_run(request_input: RequestInput):
        urllib.request.urlopen(req)

        model = ComfyDeployRunner()
-        call = model.run.spawn(request_input.input)
+        call = await model.run.spawn.aio(request_input.input)
+
+        print("call", call)

        # call = run.spawn()
-        return {"call_id": call.object_id}
+        return {"call_id": None}
    
    return {"call_id": None}

@@ -4,5 +4,7 @@ config = {
    "gpu": "T4", 
    "public_model_volume": "model-store",
    "private_model_volume": "private-model-store",
-    "pip": []
+    "pip": [],
+    "run_timeout": 60 * 5,
+    "idle_timeout": 60
 }
@@ -9,6 +9,8 @@ public:
  loras: loras
  upscale_models: upscale_models
  vae: vae
+  ipadapter: ipadapter
+

 private:
  base_path: /private_models/
@@ -21,3 +23,4 @@ private:
  loras: loras
  upscale_models: upscale_models
  vae: vae
+  ipadapter: ipadapter
@@ -2,7 +2,7 @@
  "comfyui": "d0165d819afe76bd4e6bdd710eb5f3e571b6a804",
  "git_custom_nodes": {
    "https://github.com/BennyKok/comfyui-deploy.git": {
-      "hash": "a838cb7ad425e5652c3931fbafdc886b53c48a22",
+      "hash": "df46e3a0e5ad93fa71f5d216997e376af33b2a6d",
      "disabled": false
    }
  },
@@ -1,10 +1,18 @@
+from fastapi import FastAPI, Request
+from fastapi.responses import HTMLResponse
 import modal
 from config import config
 import os
 import subprocess
 from pprint import pprint

-stub = modal.Stub()
+stub = modal.Stub(config["app_name"])
+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"]
+web_app = FastAPI()

 # 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:
@@ -21,12 +29,6 @@ def create_volumes(volume_names, paths):
 
    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"]
-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")
@@ -42,10 +44,10 @@ def download_model(volume_name, download_config):

    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
+    modified_download_url = download_url + ("&" if "?" in download_url else "?") + "token=" + civitai_key # civitai requires auth
    print('downloading', modified_download_url)

-    subprocess.run(["wget", modified_download_url , "--content-disposition", "-P", model_store_path])
+    subprocess.run(["wget", modified_download_url , "--content-disposition", "-P", model_store_path, "-nv"])
    subprocess.run(["ls", "-la", volume_base_path])
    subprocess.run(["ls", "-la", model_store_path])
    volumes[volume_base_path].commit()
@@ -56,11 +58,12 @@ def download_model(volume_name, download_config):
    print(f"finished! sending to {callback_url}")
    pprint({**status, **callback_body})

-@stub.local_entrypoint()
+@stub.function(image=image)
+# @modal.asgi_app()
 def simple_download():
    import requests
    try:
-        list(download_model.starmap([(vol_name, link) for vol_name,link in vol_name_to_links.items()]))
+        list(download_model.starmap([(vol_name, download_conf) for vol_name,download_conf 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})
@@ -71,4 +74,3 @@ def simple_download():
        requests.post(callback_url, json={**status, **callback_body})
        print(f"finished! sending to {callback_url}")
        pprint({**status, **callback_body})
-        
@@ -15,4 +15,5 @@ config = {
        "folder_path": "checkpoints",
    }, 
    "civitai_api_key": "",
+    "app_name": "",
 }
@@ -0,0 +1,12 @@
+import modal
+import requests
+from app import stub
+from config import config
+
+modal.runner.deploy_stub(stub)
+print("deployed stub")
+# web_url = stub["simple_download"].web_url
+f = modal.Function.lookup(config['app_name'], "simple_download")
+f.spawn()
+# print(f"web_url: {web_url}")
+# requests.post(web_url)
@@ -26,6 +26,8 @@ import threading
 api = None
 api_task = None
 prompt_metadata = {}
+cd_enable_log = os.environ.get('CD_ENABLE_LOG', 'false').lower() == 'true'
+cd_enable_run_log = os.environ.get('CD_ENABLE_RUN_LOG', 'false').lower() == 'true'

 def post_prompt(json_data):
    prompt_server = server.PromptServer.instance
@@ -157,7 +159,9 @@ async def websocket_handler(request):
    try:
        # Send initial state to the new client
        await send("status", { 'sid': sid }, sid)
-        await send_first_time_log(sid)
+
+        if cd_enable_log:
+            await send_first_time_log(sid)
            
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.ERROR:
@@ -236,7 +240,12 @@ class Status(Enum):
    FAILED = "failed"
    UPLOADING = "uploading"

+# Global variable to keep track of the last read line number
+last_read_line_number = 0
+
 def update_run(prompt_id, status: Status):
+    global last_read_line_number
+
    if prompt_id not in prompt_metadata:
        return

@@ -251,16 +260,50 @@ def update_run(prompt_id, status: Status):
            "run_id": prompt_id,
            "status": status.value,
        }
-        prompt_metadata[prompt_id]['status'] = status
        print(f"Status: {status.value}")

        try:
            requests.post(status_endpoint, json=body)
+
+            if cd_enable_run_log and (status == Status.SUCCESS or status == Status.FAILED):
+                try:
+                    with open(comfyui_file_path, 'r') as log_file:
+                        # log_data = log_file.read()
+                        # Move to the last read line
+                        all_log_data = log_file.read()  # Read all log data
+                        print("All log data before skipping:", all_log_data)  # Log all data before skipping
+                        log_file.seek(0)  # Reset file pointer to the beginning
+                        
+                        for _ in range(last_read_line_number):
+                            next(log_file)
+                        log_data = log_file.read()
+                        # Update the last read line number
+                        last_read_line_number += log_data.count('\n')
+                        print("last_read_line_number", last_read_line_number)
+                        print("log_data", log_data)
+                        print("log_data.count(n)", log_data.count('\n'))
+
+                        body = {
+                            "run_id": prompt_id,
+                            "log_data": [
+                                {
+                                    "logs": log_data,
+                                    # "timestamp": time.time(),
+                                }
+                            ]
+                        }
+                        requests.post(status_endpoint, json=body)
+                except Exception as log_error:
+                    print(f"Error reading log file: {log_error}")
+                
+
        except Exception as e:
            error_type = type(e).__name__
            stack_trace = traceback.format_exc().strip()
            print(f"Error occurred while updating run: {e} {stack_trace}")
-
+        finally:
+            prompt_metadata[prompt_id]['status'] = status
+            

 async def upload_file(prompt_id, filename, subfolder=None, content_type="image/png", type="output"):
    """
@@ -387,20 +430,25 @@ async def update_file_status(prompt_id, data, uploading, have_error=False, node_
            "prompt_id": prompt_id,
        })

+async def handle_upload(prompt_id, data, key, content_type_key, default_content_type):
+    items = data.get(key, [])
+    for item in items:
+        await upload_file(
+            prompt_id, 
+            item.get("filename"), 
+            subfolder=item.get("subfolder"), 
+            type=item.get("type"), 
+            content_type=item.get(content_type_key, default_content_type)
+        )
+
+
 # Upload files in the background
 async def upload_in_background(prompt_id, data, node_id=None, have_upload=True):
    try:
-        images = data.get('images', [])
-        for image in images:
-            await upload_file(prompt_id, image.get("filename"), subfolder=image.get("subfolder"), type=image.get("type"), content_type=image.get("content_type", "image/png"))
-
-        files = data.get('files', [])
-        for file in files:
-            await upload_file(prompt_id, file.get("filename"), subfolder=file.get("subfolder"), type=file.get("type"), content_type=file.get("content_type", "image/png"))
-
-        gifs = data.get('gifs', [])
-        for gif in gifs:
-            await upload_file(prompt_id, gif.get("filename"), subfolder=gif.get("subfolder"), type=gif.get("type"), content_type=gif.get("format", "image/gif"))
+        await handle_upload(prompt_id, data, 'images', "content_type", "image/png")
+        await handle_upload(prompt_id, data, 'files', "content_type", "image/png")
+        # This will also be mp4
+        await handle_upload(prompt_id, data, 'gifs', "format", "image/gif")
            
        if have_upload:
            await update_file_status(prompt_id, data, False, node_id=node_id)
@@ -441,6 +489,7 @@ prompt_server.send_json = send_json_override.__get__(prompt_server, server.Promp
 root_path = os.path.dirname(os.path.abspath(__file__))
 two_dirs_up = os.path.dirname(os.path.dirname(root_path))
 log_file_path = os.path.join(two_dirs_up, 'comfy-deploy.log')
+comfyui_file_path = os.path.join(two_dirs_up, 'comfyui.log')

 last_read_line = 0

@@ -480,4 +529,5 @@ def run_in_new_thread(coroutine):
    t.start()
    asyncio.run_coroutine_threadsafe(coroutine, new_loop)

-run_in_new_thread(watch_file_changes(log_file_path, send_logs_to_websocket))
+if cd_enable_log:
+    run_in_new_thread(watch_file_changes(log_file_path, send_logs_to_websocket))
@@ -7,45 +7,56 @@ import threading
 import logging
 from logging.handlers import RotatingFileHandler

-handler = RotatingFileHandler('comfy-deploy.log', maxBytes=500000, backupCount=5)
+# Running with export CD_ENABLE_LOG=true; python main.py

-original_stdout = sys.stdout
-original_stderr = sys.stderr
+# Check for 'cd-enable-log' flag in input arguments
+# cd_enable_log = '--cd-enable-log' in sys.argv
+cd_enable_log = os.environ.get('CD_ENABLE_LOG', 'false').lower() == 'true'

-class StreamToLogger():
-    def __init__(self, log_level):
-        self.log_level = log_level
+def setup():
+    handler = RotatingFileHandler('comfy-deploy.log', maxBytes=500000, backupCount=5)

-    def write(self, buf):
-        if (self.log_level == logging.INFO):
-            original_stdout.write(buf)
-            original_stdout.flush()
-        elif (self.log_level == logging.ERROR):
-            original_stderr.write(buf)
-            original_stderr.flush()
+    original_stdout = sys.stdout
+    original_stderr = sys.stderr

-        for line in buf.rstrip().splitlines():
-            handler.handle(
-                logging.LogRecord(
-                    name="comfy-deploy",
-                    level=self.log_level,
-                    pathname="prestartup_script.py",
-                    lineno=1,
-                    msg=line.rstrip(),
-                    args=None,
-                    exc_info=None
+    class StreamToLogger():
+        def __init__(self, log_level):
+            self.log_level = log_level
+
+        def write(self, buf):
+            if (self.log_level == logging.INFO):
+                original_stdout.write(buf)
+                original_stdout.flush()
+            elif (self.log_level == logging.ERROR):
+                original_stderr.write(buf)
+                original_stderr.flush()
+
+            for line in buf.rstrip().splitlines():
+                handler.handle(
+                    logging.LogRecord(
+                        name="comfy-deploy",
+                        level=self.log_level,
+                        pathname="prestartup_script.py",
+                        lineno=1,
+                        msg=line.rstrip(),
+                        args=None,
+                        exc_info=None
+                    )
                )
-            )

-    def flush(self):
-        if (self.log_level == logging.INFO):
-            original_stdout.flush()
-        elif (self.log_level == logging.ERROR):
-            original_stderr.flush()
+        def flush(self):
+            if (self.log_level == logging.INFO):
+                original_stdout.flush()
+            elif (self.log_level == logging.ERROR):
+                original_stderr.flush()

-# Redirect stdout and stderr to the logger
-sys.stdout = StreamToLogger(logging.INFO)
-sys.stderr = StreamToLogger(logging.ERROR)
+    # Redirect stdout and stderr to the logger
+    sys.stdout = StreamToLogger(logging.INFO)
+    sys.stderr = StreamToLogger(logging.ERROR)
+
+if cd_enable_log:
+   print("** Comfy Deploy logging enabled")
+   setup() 

 try:
    # Get the absolute path of the script's directory
@@ -50,11 +50,22 @@ const ext = {
      })
        .then(async (res) => {
          const data = await res.json();
-          const { workflow, error } = data;
+          const { workflow, workflow_id, error } = data;
          if (error) {
            infoDialog.showMessage("Unable to load this workflow", error);
            return;
          }
+
+          // Adding a delay to wait for the intial graph to load
+          await new Promise((resolve) => setTimeout(resolve, 2000));
+
+          workflow?.nodes.forEach((x) => {
+            if (x?.type === "ComfyDeploy") {
+              x.widgets_values[1] = workflow_id;
+              // x.widgets_values[2] = workflow_version.version;
+            }
+          });
+
          /** @type {LGraph} */
          app.loadGraphData(workflow);
        })
@@ -682,16 +693,19 @@ export class ConfigDialog extends ComfyDialog {
    </label>
      <label style="color: white; width: 100%;">
        Endpoint:
-        <input id="endpoint" style="margin-top: 8px; width: 100%; height:40px; box-sizing: border-box; padding: 0px 6px;" type="text" value="${data.endpoint
-      }">
+        <input id="endpoint" style="margin-top: 8px; width: 100%; height:40px; box-sizing: border-box; padding: 0px 6px;" type="text" value="${
+          data.endpoint
+        }">
      </label>
      <label style="color: white;">
        API Key: ${data.displayName ?? ""}
-        <input id="apiKey" style="margin-top: 8px; width: 100%; height:40px; box-sizing: border-box; padding: 0px 6px;" type="password" value="${data.apiKey
-      }">
+        <input id="apiKey" style="margin-top: 8px; width: 100%; height:40px; box-sizing: border-box; padding: 0px 6px;" type="password" value="${
+          data.apiKey
+        }">
        <button id="loginButton" style="margin-top: 8px; width: 100%; height:40px; box-sizing: border-box; padding: 0px 6px;">
-          ${data.apiKey ? "Re-login with ComfyDeploy" : "Login with ComfyDeploy"
-      }
+          ${
+            data.apiKey ? "Re-login with ComfyDeploy" : "Login with ComfyDeploy"
+          }
        </button>
      </label>
      </div>
@@ -0,0 +1 @@
+ALTER TYPE "model_upload_type" ADD VALUE 'download_url';
@@ -0,0 +1 @@
+ALTER TYPE "model_upload_type" ADD VALUE 'download-url';
@@ -0,0 +1,2 @@
+ALTER TYPE "workflow_run_status" ADD VALUE 'timeout';--> statement-breakpoint
+ALTER TABLE "comfyui_deploy"."workflow_runs" ADD COLUMN "run_log" text;
@@ -0,0 +1,2 @@
+ALTER TABLE "comfyui_deploy"."workflow_runs" DROP COLUMN "run_log";
+ALTER TABLE "comfyui_deploy"."workflow_runs" ADD COLUMN "run_log" jsonb;
@@ -323,6 +323,34 @@
      "when": 1706336448134,
      "tag": "0045_careful_cerise",
      "breakpoints": true
+    },
+    {
+      "idx": 46,
+      "version": "5",
+      "when": 1706383154642,
+      "tag": "0046_complex_mentallo",
+      "breakpoints": true
+    },
+    {
+      "idx": 47,
+      "version": "5",
+      "when": 1706384528895,
+      "tag": "0047_gifted_starbolt",
+      "breakpoints": true
+    },
+    {
+      "idx": 48,
+      "version": "5",
+      "when": 1706600255919,
+      "tag": "0048_dear_korath",
+      "breakpoints": true
+    },
+    {
+      "idx": 49,
+      "version": "5",
+      "when": 1706631744127,
+      "tag": "0049_sweet_hex",
+      "breakpoints": true
    }
  ]
-}
+}
@@ -36,6 +36,7 @@ export async function GET(req: Request) {
      orgId: orgId ?? null,
      plan: plan,
    },
+    allow_promotion_codes: true,
    client_reference_id: orgId ?? userId,
    customer_email: user.emailAddresses[0].emailAddress,
    mode: "subscription",
@@ -2,10 +2,8 @@ import { parseDataSafe } from "../../../../lib/parseDataSafe";
 import { db } from "@/db/db";
 import {
  WorkflowRunStatusSchema,
-  userUsageTable,
  workflowRunOutputs,
  workflowRunsTable,
-  workflowTable,
 } from "@/db/schema";
 import { getCurrentPlan } from "@/server/getCurrentPlan";
 import { stripe } from "@/server/stripe";
@@ -18,6 +16,7 @@ const Request = z.object({
  status: WorkflowRunStatusSchema.optional(),
  time: z.coerce.date().optional(),
  output_data: z.any().optional(),
+  log_data: z.any().optional(),
 });

 export async function POST(request: Request) {
@@ -26,7 +25,26 @@ export async function POST(request: Request) {

    if (!data || error) return error;

-    const { run_id, status, time, output_data } = data;
+    const { run_id, status, time, output_data, log_data } = data;
+
+    if (log_data) {
+      // It successfully started, update the started_at time
+      await db
+        .update(workflowRunsTable)
+        .set({
+          run_log: log_data,
+        })
+        .where(eq(workflowRunsTable.id, run_id));
+
+      return NextResponse.json(
+        {
+          message: "success",
+        },
+        {
+          status: 200,
+        },
+      );
+    }

    if (status == "started" && time != undefined) {
      // It successfully started, update the started_at time
@@ -48,6 +66,9 @@ export async function POST(request: Request) {
        .where(eq(workflowRunsTable.id, run_id));
    }

+    const ended =
+      status === "success" || status === "failed" || status === "timeout";
+
    if (output_data) {
      const workflow_run_output = await db.insert(workflowRunOutputs).values({
        run_id: run_id,
@@ -58,8 +79,7 @@ export async function POST(request: Request) {
        .update(workflowRunsTable)
        .set({
          status: status,
-          ended_at:
-            status === "success" || status === "failed" ? new Date() : null,
+          ended_at: ended ? new Date() : null,
        })
        .where(eq(workflowRunsTable.id, run_id))
        .returning();
@@ -67,10 +87,7 @@ export async function POST(request: Request) {
      // Need to filter out only comfy deploy serverless
      // Also multiply with the gpu selection
      if (workflow_run.machine_type == "comfy-deploy-serverless") {
-        if (
-          (status === "success" || status === "failed") &&
-          workflow_run.user_id
-        ) {
+        if (ended && workflow_run.user_id) {
          const sub = await getCurrentPlan({
            user_id: workflow_run.user_id,
            org_id: workflow_run.org_id,
@@ -91,12 +108,16 @@ export async function POST(request: Request) {
                durationInSec *= 4;
                break;
            }
-            await stripe.subscriptionItems.createUsageRecord(
-              sub.subscription_item_api_id,
-              {
-                quantity: durationInSec,
-              },
-            );
+            try {
+              await stripe.subscriptionItems.createUsageRecord(
+                sub.subscription_item_api_id,
+                {
+                  quantity: durationInSec,
+                },
+              );
+            } catch (error) {
+              console.log(error);
+            }
          }
        }
      }
@@ -35,7 +35,10 @@ export default async function Page({
            />
          )}
          {machine.status !== "building" && machine.build_log && (
-            <LogsViewer logs={JSON.parse(machine.build_log)} />
+            <LogsViewer
+              logs={JSON.parse(machine.build_log)}
+              className="h-full max-h-[600px]"
+            />
          )}
        </CardContent>
      </Card>
@@ -1,7 +1,9 @@
 import { AccessType } from "../../../lib/AccessType";
 import { MachineList } from "@/components/MachineList";
+import { SubscriptionProvider } from "@/components/useCurrentPlan";
 import { db } from "@/db/db";
 import { machinesTable } from "@/db/schema";
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
 import { auth } from "@clerk/nextjs";
 import { clerkClient } from "@clerk/nextjs/server";
 import { desc, eq, isNull, and } from "drizzle-orm";
@@ -19,6 +21,8 @@ async function MachineListServer() {

  const user = await clerkClient.users.getUser(userId);

+  const sub = await getCurrentPlanWithAuth();
+
  const machines = await db.query.machinesTable.findMany({
    orderBy: desc(machinesTable.updated_at),
    where:
@@ -29,11 +33,12 @@ async function MachineListServer() {

  return (
    <div className="w-full">
-      {/* <div>Machines</div> */}
-      <MachineList
-        data={machines}
-        userMetadata={AccessType.parse(user.privateMetadata ?? {})}
-      />
+      <SubscriptionProvider sub={sub}>
+        <MachineList
+          data={machines}
+          userMetadata={AccessType.parse(user.privateMetadata ?? {})}
+        />
+      </SubscriptionProvider>
    </div>
  );
 }
@@ -16,11 +16,19 @@ const freeTierSeconds = 30000;
 export default async function Home() {
  const sub = await getCurrentPlanWithAuth();

-  const data = sub?.subscription_item_api_id
-    ? await stripe.subscriptionItems.listUsageRecordSummaries(
-        sub?.subscription_item_api_id,
-      )
-    : null;
+  let data: Awaited<
+    ReturnType<typeof stripe.subscriptionItems.listUsageRecordSummaries>
+  > | null = null;
+
+  try {
+    data = sub?.subscription_item_api_id
+      ? await stripe.subscriptionItems.listUsageRecordSummaries(
+          sub?.subscription_item_api_id,
+        )
+      : null;
+  } catch (e) {
+    console.error(e);
+  }

  return (
    <div className="mt-4 flex items-center justify-center">
@@ -1,5 +1,5 @@
 import { LoadingWrapper } from "@/components/LoadingWrapper";
-import { DeploymentsTable } from "@/components/RunsTable";
+import { DeploymentsTable } from "@/components/DeploymentsTable";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";

 export default async function Page({
@@ -13,12 +13,12 @@ export default async function Page({
  const workflow_id = params.workflow_id;

  return (
-    <Card className="w-full h-fit min-w-0">
-      <CardHeader className="relative">
+    <Card className="w-full h-fit min-w-0 relative">
+      <CardHeader>
        <CardTitle>Run</CardTitle>
-        <div className="absolute right-6 top-6">
+        {/* <div className="absolute right-6 top-6">
          <RouteRefresher interval={5000} autoRefresh={false} />
-        </div>
+        </div> */}
      </CardHeader>

      <CardContent>
@@ -8,8 +8,8 @@ import {
  OpenEditButton,
  RunWorkflowButton,
  VersionSelect,
-  ViewWorkflowDetailsButton,
 } from "@/components/VersionSelect";
+import { ViewWorkflowDetailsButton } from "@/components/ViewWorkflowDetailsButton";
 import {
  Card,
  CardContent,
@@ -1,12 +1,12 @@
 export const metadata = {
  title: 'Quickstart',
  description:
-    'This guide will get you all set up and ready to use the Protocol API. We’ll cover how to get started an API client and how to make your first API request.',
+    'This guide will get you all set up and ready to use Comfy Deploy. We’ll cover how to get started an API client and how to make your first API request.',
 }

-# Getting stated
+# Getting started

-Install Comfy Deploy's plugin on your local machine to get started with deploying workflow.
+Install Comfy Deploy's plugin on your local ComfyUI to get started with deploying workflow.

 <CodeGroup>

@@ -0,0 +1,58 @@
+export const metadata = {
+  title: 'Comfy Deploy Video tutorials',
+  description:
+    'Get started with API integration to run any deploy ComfyUI workflow',
+}
+
+**This is a collection of video tutorials on questions that have come up in our [discord channel](https://discord.gg/Zrar7yP6MJ).**
+
+# Comfy Deploy workflow walkthrough (comprehensive)
+
+Local workflow, to production API
+- Set inputs with Comfy Deploy `External Text` nodes
+- Deploy the workflow 
+- Create a machine
+- Deploy an API endpoint for the workflow
+- Test the workflow in Comfy Deploy
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/comfydeploy_base-sd-setup.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+
+
+# Test with Docs
+Shows how you can use the <a href='/docs/endpoints'>endpoints docs</a> to test out the API!
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/comfydeploy_interactive_docs.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+## Intermediate IPAdaptor example run 
+Shows off a semi-complex workflow working on Comfy Deploy
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/ipadapters-test.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+## Animate diff on Comfy Deploy
+Shows off animate diff working on Comfy Deploy
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/comfydeploy-animediff.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+## Cloning workflows
+What does cloning a workflow look like
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/comfydeploy-clone-template.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+## [Install custom nodes from any github repo (Loom)](https://www.loom.com/share/c2025c5060e348839ade6a3a9d96441d?sid=6a57f779-0209-4341-a5fd-f99769cdb162)
+- This is only neccessary if the custom node is not in the ComfyUI Manager node list
@@ -13,10 +13,7 @@ export async function CodeBlock(props: {

  return (
    <div className="relative w-full text-sm">
-      {/* max-w-[calc(32rem-1.5rem-1.5rem)] */}
-      {/* <div className=""> */}
      <p
-        // tabIndex={1}
        className="[&>pre]:p-4 rounded-lg max-h-96 overflow-auto w-full"
        style={{
          overflowWrap: "break-word",
@@ -28,7 +25,6 @@ export async function CodeBlock(props: {
          }),
        }}
      />
-      {/* </div> */}
      <CopyButton className="absolute right-2 top-2" text={props.code} />
    </div>
  );
@@ -0,0 +1,35 @@
+"use client";
+
+import { CopyButton } from "@/components/CopyButton";
+import type { StringLiteralUnion } from "shikiji";
+import useSWR from "swr";
+import { highlight } from "../server/highlight";
+
+export function CodeBlockClient({
+  code,
+  lang,
+}: {
+  code: string;
+  lang: StringLiteralUnion<string>;
+}) {
+  const { data } = useSWR(code, async () => {
+    return highlight(code.trim(), lang);
+  });
+
+  return (
+    <div className="relative w-full text-sm">
+      {data && (
+        <p
+          className="[&>pre]:p-4 rounded-lg max-h-96 overflow-auto w-full"
+          style={{
+            overflowWrap: "break-word",
+          }}
+          dangerouslySetInnerHTML={{
+            __html: data,
+          }}
+        />
+      )}
+      <CopyButton className="absolute right-2 top-2" text={code} />
+    </div>
+  );
+}
@@ -7,10 +7,12 @@ import { toast } from "sonner";

 export function CopyButton({
  className,
+  children,
  ...props
 }: {
  text: string;
  className?: string;
+  children?: React.ReactNode;
 }) {
  return (
    <Button
@@ -21,7 +23,7 @@ export function CopyButton({
      }}
      className={cn(" p-2 min-h-0 aspect-square", className)}
    >
-      <Copy size={14} />
+      {children} <Copy size={14} />
    </Button>
  );
 }
@@ -0,0 +1,24 @@
+"use client";
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
+import * as React from "react";
+import { createContext } from "react";
+
+export type CurrentPlanContextType = Awaited<
+  ReturnType<typeof getCurrentPlanWithAuth>
+>;
+export const CurrentPlanContext = createContext<
+  CurrentPlanContextType | undefined
+>(undefined);
+export function SubscriptionProvider({
+  sub,
+  children,
+}: {
+  sub: CurrentPlanContextType;
+  children: React.ReactNode;
+}) {
+  return (
+    <CurrentPlanContext.Provider value={sub}>
+      {children}
+    </CurrentPlanContext.Provider>
+  );
+}
@@ -69,13 +69,16 @@ const client = new ComfyDeployClient({
 `;

 const jsClientCreateRunTemplate = `
-const { run_id } = await client.run("<ID>", {
+const { run_id } = await client.run({
+  deployment_id: "<ID>",
  inputs: {}
 });
 `;

 const jsClientCreateRunNoInputsTemplate = `
-const { run_id } = await client.run("<ID>");
+const { run_id } = await client.run({
+  deployment_id: "<ID>"
+});
 `;

 const clientTemplate_checkStatus = `
@@ -118,15 +121,9 @@ export function DeploymentDisplay({
            </TabsList>
            <TabsContent className="flex flex-col gap-2 !mt-0" value="client">
              <div>
-                Copy and paste the ComfyDeployClient form&nbsp;
-                <a
-                  href="https://github.com/BennyKok/comfyui-deploy-next-example/blob/main/src/lib/comfy-deploy.ts"
-                  className="text-blue-500 hover:underline"
-                  target="_blank"
-                  rel="noreferrer"
-                >
-                  here
-                </a>
+                Install the node comfydeploy SDK
+                <CodeBlock lang="bash" code={`npm i comfydeploy`} />
+                Initialize your client
              </div>
              <CodeBlock
                lang="js"
@@ -0,0 +1,41 @@
+import {
+  Table,
+  TableBody,
+  TableCaption,
+  TableHead,
+  TableHeader,
+  TableRow,
+} from "@/components/ui/table";
+import { headers } from "next/headers";
+import { findAllDeployments } from "../server/findAllRuns";
+import { DeploymentDisplay } from "./DeploymentDisplay";
+
+export async function DeploymentsTable(props: { workflow_id: string }) {
+  const allRuns = await findAllDeployments(props.workflow_id);
+
+  const headersList = headers();
+  const host = headersList.get("host") || "";
+  const protocol = headersList.get("x-forwarded-proto") || "";
+  const domain = `${protocol}://${host}`;
+
+  return (
+    <div className="overflow-auto h-fit w-full">
+      <Table className="">
+        <TableCaption>A list of your deployments</TableCaption>
+        <TableHeader className="bg-background top-0 sticky">
+          <TableRow>
+            <TableHead className=" w-[100px]">Environment</TableHead>
+            <TableHead className=" w-[100px]">Version</TableHead>
+            <TableHead className="">Machine</TableHead>
+            <TableHead className=" text-right">Updated At</TableHead>
+          </TableRow>
+        </TableHeader>
+        <TableBody>
+          {allRuns.map((run) => (
+            <DeploymentDisplay deployment={run} key={run.id} domain={domain} />
+          ))}
+        </TableBody>
+      </Table>
+    </div>
+  );
+}
@@ -23,14 +23,42 @@ import * as React from "react";
 import { useState } from "react";
 import type { UnknownKeysParam, ZodObject, ZodRawShape, z } from "zod";

+type ContextType = [Partial<any>, React.Dispatch<React.SetStateAction<any>>];
+const AutoFormValueContext = React.createContext<ContextType | null>(null);
+
+function AutoFormValueProvider<Z extends ZodObject<any, any>>({
+  children,
+  value,
+}: {
+  children: React.ReactNode;
+  value: ContextType;
+}) {
+  return (
+    <AutoFormValueContext.Provider value={value}>
+      {children}
+    </AutoFormValueContext.Provider>
+  );
+}
+
+export function useAutoFormValueContext<Z extends ZodObject<any, any>>() {
+  const context = React.useContext(AutoFormValueContext);
+
+  // if (!context) {
+  //   throw new Error("useInsertModal must be used within a InsertModalProvider");
+  // }
+
+  return context;
+}
+
 export function InsertModal<
  K extends ZodRawShape,
  Y extends UnknownKeysParam,
-  Z extends ZodObject<K, Y>
+  Z extends ZodObject<K, Y>,
 >(props: {
  tooltip?: string;
  disabled?: boolean;
-  title: string;
+  title: React.ReactNode;
+  buttonTitle?: React.ReactNode;
  description: string;
  dialogClassName?: string;
  serverAction: (data: z.infer<Z>) => Promise<unknown>;
@@ -40,72 +68,80 @@ export function InsertModal<
  const [open, setOpen] = React.useState(false);
  const [isLoading, setIsLoading] = React.useState(false);

+  const [values, setValues] = useState<Partial<z.infer<Z>>>({});
+
  return (
-    <Dialog open={open} onOpenChange={setOpen}>
-      {/* <DialogTrigger disabled={props.disabled}> */}
-      {props.tooltip ? (
-        <Tooltip>
-          <TooltipTrigger asChild>
-            <Button
-              variant="default"
-              className={props.disabled ? "opacity-50" : ""}
-              onClick={() => {
-                if (props.disabled) return;
-                setOpen(true);
-              }}
-            >
-              {props.title}
-            </Button>
-          </TooltipTrigger>
-          <TooltipContent>
-            <p>{props.tooltip}</p>
-          </TooltipContent>
-        </Tooltip>
-      ) : (
-        <Button
-          variant="default"
-          disabled={props.disabled}
-          onClick={() => {
-            setOpen(true);
-          }}
+    <AutoFormValueProvider value={[values, setValues]}>
+      <Dialog open={open} onOpenChange={setOpen}>
+        {/* <DialogTrigger disabled={props.disabled}> */}
+        {props.tooltip ? (
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                variant="default"
+                className={props.disabled ? "opacity-50" : ""}
+                onClick={() => {
+                  if (props.disabled) return;
+                  setOpen(true);
+                }}
+              >
+                {props.buttonTitle ?? props.title}
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>
+              <p>{props.tooltip}</p>
+            </TooltipContent>
+          </Tooltip>
+        ) : (
+          <Button
+            variant="default"
+            disabled={props.disabled}
+            onClick={() => {
+              setOpen(true);
+            }}
+          >
+            {props.title}
+          </Button>
+        )}
+        {/* </DialogTrigger> */}
+        <DialogContent
+          className={cn("sm:max-w-[425px]", props.dialogClassName)}
        >
-          {props.title}
-        </Button>
-      )}
-      {/* </DialogTrigger> */}
-      <DialogContent className={cn("sm:max-w-[425px]", props.dialogClassName)}>
-        <DialogHeader>
-          <DialogTitle>{props.title}</DialogTitle>
-          <DialogDescription>{props.description}</DialogDescription>
-        </DialogHeader>
-        {/* <ScrollArea> */}
-        <AutoForm
-          fieldConfig={props.fieldConfig}
-          formSchema={props.formSchema}
-          onSubmit={async (data) => {
-            setIsLoading(true);
-            await callServerPromise(props.serverAction(data));
-            setIsLoading(false);
-            setOpen(false);
-          }}
-        >
-          <div className="flex justify-end">
-            <AutoFormSubmit>
-              Save Changes
-              {isLoading && <LoadingIcon />}
-            </AutoFormSubmit>
-          </div>
-        </AutoForm>
-        {/* </ScrollArea> */}
-      </DialogContent>
-    </Dialog>
+          <DialogHeader>
+            <DialogTitle>{props.title}</DialogTitle>
+            <DialogDescription>{props.description}</DialogDescription>
+          </DialogHeader>
+          {/* <ScrollArea> */}
+          <AutoForm
+            values={values}
+            onValuesChange={setValues}
+            fieldConfig={props.fieldConfig}
+            formSchema={props.formSchema}
+            onSubmit={async (data) => {
+              setIsLoading(true);
+              await callServerPromise(props.serverAction(data));
+              setIsLoading(false);
+              setOpen(false);
+            }}
+          >
+            <div className="flex justify-end">
+              <AutoFormSubmit>
+                Save Changes
+                {isLoading && <LoadingIcon />}
+              </AutoFormSubmit>
+            </div>
+          </AutoForm>
+          {/* </ScrollArea> */}
+        </DialogContent>
+      </Dialog>
+    </AutoFormValueProvider>
  );
 }

 export function UpdateModal<
  K extends ZodRawShape,
  Y extends UnknownKeysParam,
-  Z extends ZodObject<K, Y>
+  Z extends ZodObject<K, Y>,
 >(props: {
  open?: boolean;
  setOpen?: (open: boolean) => void;
@@ -118,7 +154,7 @@ export function UpdateModal<
  serverAction: (
    data: z.infer<Z> & {
      id: string;
-    }
+    },
  ) => Promise<unknown>;
  formSchema: Z;
  fieldConfig?: FieldConfig<z.infer<Z>>;
@@ -137,49 +173,53 @@ export function UpdateModal<
  }, [props.data]);

  return (
-    <Dialog open={open} onOpenChange={setOpen}>
-      {props.trigger ?? (
-        <DialogTrigger
-          className="appearance-none hover:cursor-pointer"
-          asChild
-          onClick={() => {
-            setOpen(true);
-          }}
+    <AutoFormValueProvider value={[values, setValues]}>
+      <Dialog open={open} onOpenChange={setOpen}>
+        {props.trigger ?? (
+          <DialogTrigger
+            className="appearance-none hover:cursor-pointer"
+            asChild
+            onClick={() => {
+              setOpen(true);
+            }}
+          >
+            {props.trigger}
+          </DialogTrigger>
+        )}
+        <DialogContent
+          className={cn("sm:max-w-[425px]", props.dialogClassName)}
        >
-          {props.trigger}
-        </DialogTrigger>
-      )}
-      <DialogContent className={cn("sm:max-w-[425px]", props.dialogClassName)}>
-        <DialogHeader>
-          <DialogTitle>{props.title}</DialogTitle>
-          <DialogDescription>{props.description}</DialogDescription>
-        </DialogHeader>
-        <AutoForm
-          values={values}
-          onValuesChange={setValues}
-          fieldConfig={props.fieldConfig}
-          formSchema={props.formSchema}
-          onSubmit={async (data) => {
-            setIsLoading(true);
-            await callServerPromise(
-              props.serverAction({
-                ...data,
-                id: props.data.id,
-              })
-            );
-            setIsLoading(false);
-            setOpen(false);
-          }}
-        >
-          <div className="flex justify-end flex-wrap gap-2">
-            {props.extraButtons}
-            <AutoFormSubmit>
-              Save Changes
-              {isLoading && <LoadingIcon />}
-            </AutoFormSubmit>
-          </div>
-        </AutoForm>
-      </DialogContent>
-    </Dialog>
+          <DialogHeader>
+            <DialogTitle>{props.title}</DialogTitle>
+            <DialogDescription>{props.description}</DialogDescription>
+          </DialogHeader>
+          <AutoForm
+            values={values}
+            onValuesChange={setValues}
+            fieldConfig={props.fieldConfig}
+            formSchema={props.formSchema}
+            onSubmit={async (data) => {
+              setIsLoading(true);
+              await callServerPromise(
+                props.serverAction({
+                  ...data,
+                  id: props.data.id,
+                }),
+              );
+              setIsLoading(false);
+              setOpen(false);
+            }}
+          >
+            <div className="flex justify-end flex-wrap gap-2">
+              {props.extraButtons}
+              <AutoFormSubmit>
+                Save Changes
+                {isLoading && <LoadingIcon />}
+              </AutoFormSubmit>
+            </div>
+          </AutoForm>
+        </DialogContent>
+      </Dialog>
+    </AutoFormValueProvider>
  );
 }
@@ -16,7 +16,7 @@ export function LiveStatus({
    (state) =>
      state.data
        .filter((x) => x.id === run.id)
-        .sort((a, b) => b.timestamp - a.timestamp)?.[0]
+        .sort((a, b) => b.timestamp - a.timestamp)?.[0],
  );

  let status = run.status;
@@ -51,7 +51,9 @@ export function LiveStatus({
    <>
      <TableCell>
        {data && status != "success"
-          ? `${data.json.event} - ${data.json.data.node}`
+          ? `${data.json.event}${
+              data.json.data.node ? " - " + data.json.data.node : ""
+            }`
          : "-"}
      </TableCell>
      <TableCell className="truncate text-right">
@@ -1,17 +1,31 @@
 "use client";

 import React, { useEffect, useRef } from "react";
+import { cn } from "@/lib/utils";
+import { toast } from "sonner";

 export type LogsType = {
  machine_id?: string;
  logs: string;
-  timestamp: number;
+  timestamp?: number;
 }[];

-export function LogsViewer({ logs }: { logs: LogsType }) {
+export function LogsViewer({
+  logs,
+  hideTimestamp,
+  className,
+  stickToBottom = true,
+}: {
+  logs: LogsType;
+  hideTimestamp?: boolean;
+  className?: string;
+  stickToBottom?: boolean;
+}) {
  const container = useRef<HTMLDivElement | null>(null);

  useEffect(() => {
+    if (!stickToBottom) return;
+
    // console.log(logs.length, container.current);
    if (container.current) {
      const scrollHeight = container.current.scrollHeight;
@@ -21,11 +35,12 @@ export function LogsViewer({ logs }: { logs: LogsType }) {
        behavior: "smooth",
      });
    }
-  }, [logs.length]);
+  }, [logs.length, stickToBottom]);

  return (
    <div
      ref={(ref) => {
+        if (!stickToBottom) return;
        if (!container.current && ref) {
          const scrollHeight = ref.scrollHeight;

@@ -36,10 +51,31 @@ export function LogsViewer({ logs }: { logs: LogsType }) {
        }
        container.current = ref;
      }}
-      className="flex flex-col text-xs p-2 overflow-y-scroll max-h-[400px] whitespace-break-spaces"
+      className={cn(
+        "h-full w-full flex flex-col text-xs p-2 overflow-y-scroll whitespace-break-spaces",
+        className,
+      )}
    >
      {logs.map((x, i) => (
-        <div key={i}>{x.logs}</div>
+        <div
+          key={i}
+          className="hover:bg-gray-100 flex flex-row items-center gap-2"
+          onClick={() => {
+            toast.success("Copied to clipboard");
+            navigator.clipboard.writeText(x.logs);
+          }}
+        >
+          {!hideTimestamp && x.timestamp != undefined && (
+            <>
+              <span className="w-[150px] flex-shrink-0">
+                {new Date(x.timestamp * 1000).toLocaleString()}
+              </span>
+              <div className="h-full w-[1px] bg-stone-400 flex-shrink-0"></div>
+            </>
+          )}
+          {/* Display timestamp */}
+          <div>{x.logs}</div>
+        </div>
      ))}
    </div>
  );
@@ -16,7 +16,7 @@ import {
  AlertDialogHeader,
  AlertDialogTitle,
  AlertDialogTrigger,
-} from "@/components/ui/alert-dialog"
+} from "@/components/ui/alert-dialog";
 import { useRouter } from "next/navigation";

 export function MachineBuildLog({
@@ -41,7 +41,7 @@ export function MachineBuildLog({
      reconnectAttempts: 20,
      reconnectInterval: 1000,
      queryParams: query,
-    }
+    },
  );

  const connectionStatus = getConnectionStatus(readyState);
@@ -57,59 +57,70 @@ export function MachineBuildLog({
      setLogs((logs) => [...(logs ?? []), message.data]);
    } else if (message?.event === "FINISHED") {
      setFinished(true);
-      setStatus(message.data.status)
+      setStatus(message.data.status);
    }
  }, [lastMessage]);

-  const router = useRouter()
+  const router = useRouter();

  return (
    <div>
      {connectionStatus}
-      <LogsViewer logs={logs} />
+      <LogsViewer logs={logs} className="h-full max-h-[600px]" />

      <AlertDialog open={finished}>
        <AlertDialogContent>
-          {
-            status == "succuss" ? (
-              <>
-                <AlertDialogHeader>
-                  <AlertDialogTitle>Machine Built</AlertDialogTitle>
-                  <AlertDialogDescription>
-                    Your machine is built, you can now integrate your API, or directly run to check this machines.
-                  </AlertDialogDescription>
-                </AlertDialogHeader>
-                <AlertDialogFooter>
-                  <AlertDialogAction onClick={() => {
-                    router.push("/workflows")
-                  }}>See Workflows</AlertDialogAction>
-                  <AlertDialogAction onClick={() => {
-                    router.push("/machines")
-                  }}>See All Machines</AlertDialogAction>
-                </AlertDialogFooter></>
-            ) : (
-              <>
-                <AlertDialogHeader>
-                  <AlertDialogTitle>Machine Failed</AlertDialogTitle>
-                  <AlertDialogDescription>
-                    Something went wrong with the machine build, please check the log.
-                    Possible cause could be conflits with custom nodes, build got stuck, timeout, or too many custom nodes installed.
-                    Please attempt a rebuild or remove some of the custom nodes.
-                  </AlertDialogDescription>
-                </AlertDialogHeader>
-                <AlertDialogFooter>
-                  <AlertDialogCancel>See logs</AlertDialogCancel>
-                  <AlertDialogAction onClick={() => {
-                    router.push("/machines")
-                  }}>Back to machines</AlertDialogAction>
-                </AlertDialogFooter></>
-            )
-          }
-
-
+          {status == "succuss" ? (
+            <>
+              <AlertDialogHeader>
+                <AlertDialogTitle>Machine Built</AlertDialogTitle>
+                <AlertDialogDescription>
+                  Your machine is built, you can now integrate your API, or
+                  directly run to check this machines.
+                </AlertDialogDescription>
+              </AlertDialogHeader>
+              <AlertDialogFooter>
+                <AlertDialogAction
+                  onClick={() => {
+                    router.push("/workflows");
+                  }}
+                >
+                  See Workflows
+                </AlertDialogAction>
+                <AlertDialogAction
+                  onClick={() => {
+                    router.push("/machines");
+                  }}
+                >
+                  See All Machines
+                </AlertDialogAction>
+              </AlertDialogFooter>
+            </>
+          ) : (
+            <>
+              <AlertDialogHeader>
+                <AlertDialogTitle>Machine Failed</AlertDialogTitle>
+                <AlertDialogDescription>
+                  Something went wrong with the machine build, please check the
+                  log. Possible cause could be conflits with custom nodes, build
+                  got stuck, timeout, or too many custom nodes installed. Please
+                  attempt a rebuild or remove some of the custom nodes.
+                </AlertDialogDescription>
+              </AlertDialogHeader>
+              <AlertDialogFooter>
+                <AlertDialogCancel>See logs</AlertDialogCancel>
+                <AlertDialogAction
+                  onClick={() => {
+                    router.push("/machines");
+                  }}
+                >
+                  Back to machines
+                </AlertDialogAction>
+              </AlertDialogFooter>
+            </>
+          )}
        </AlertDialogContent>
      </AlertDialog>
-
    </div>
  );
 }
@@ -41,6 +41,7 @@ import {
  updateMachine,
 } from "@/server/curdMachine";
 import { editWorkflowOnMachine } from "@/server/editWorkflowOnMachine";
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
 import type {
  ColumnDef,
  ColumnFiltersState,
@@ -55,7 +56,7 @@ import {
  getSortedRowModel,
  useReactTable,
 } from "@tanstack/react-table";
-import { ArrowUpDown, MoreHorizontal } from "lucide-react";
+import { ArrowUpDown, Lock, MoreHorizontal, Plus } from "lucide-react";
 import * as React from "react";
 import { useState } from "react";
 import { toast } from "sonner";
@@ -148,9 +149,16 @@ export const columns: ColumnDef<Machine>[] = [
    header: () => <div className="text-left">Type</div>,
    cell: ({ row }) => {
      return (
-        <div className="text-left font-medium truncate">
+        <Badge
+          className="text-left font-medium truncate"
+          variant={
+            row.original.type == "comfy-deploy-serverless"
+              ? "success"
+              : "outline"
+          }
+        >
          {row.original.type}
-        </div>
+        </Badge>
      );
    },
  },
@@ -182,6 +190,7 @@ export const columns: ColumnDef<Machine>[] = [
    cell: ({ row }) => {
      const machine = row.original;
      const [open, setOpen] = useState(false);
+      const sub = useCurrentPlan();

      return (
        <DropdownMenu>
@@ -290,7 +299,10 @@ export const columns: ColumnDef<Machine>[] = [
                  fieldType: "models",
                },
                gpu: {
-                  inputProps: {},
+                  fieldType: "gpuPicker",
+                  inputProps: {
+                    sub: sub,
+                  },
                },
              }}
            />
@@ -318,6 +330,8 @@ export const columns: ColumnDef<Machine>[] = [
  },
 ];

+import { useCurrentPlan } from "./useCurrentPlan";
+
 export function MachineList({
  data,
  userMetadata,
@@ -333,6 +347,8 @@ export function MachineList({
    React.useState<VisibilityState>({});
  const [rowSelection, setRowSelection] = React.useState({});

+  const sub = useCurrentPlan();
+
  const table = useReactTable({
    data,
    columns,
@@ -352,6 +368,21 @@ export function MachineList({
    },
  });

+  let machineMaxCount = 2;
+
+  // Temp fixes for machine count
+  if (userMetadata.betaFeaturesAccess) machineMaxCount = 5;
+
+  if (sub?.plan == "pro") {
+    machineMaxCount = 10;
+  } else if (sub?.plan == "enterprise") {
+    machineMaxCount = 99;
+  }
+
+  const locked =
+    data.some((machine) => machine.type === "modal-serverless") &&
+    data.length >= machineMaxCount;
+
  return (
    <div className="w-full">
      <div className="flex items-center py-4">
@@ -366,17 +397,18 @@ export function MachineList({
        <div className="ml-auto flex gap-2">
          <InsertModal
            dialogClassName="sm:max-w-[600px]"
-            disabled={
-              data.some(
-                (machine) => machine.type === "comfy-deploy-serverless",
-              ) && !userMetadata.betaFeaturesAccess
-            }
+            disabled={locked}
            tooltip={
-              data.some((machine) => machine.type === "comfy-deploy-serverless")
-                ? "Only one hosted machine at preview stage"
-                : undefined
+              locked
+                ? `Max ${machineMaxCount} ComfyUI machine for your account, upgrade to unlock more cnfiguration.`
+                : `Max ${machineMaxCount} ComfyUI machine for your account`
            }
-            title="New Machine"
+            buttonTitle={
+              <>
+                New Machine {locked ? <Lock size={14} /> : <Plus size={14} />}
+              </>
+            }
+            title={"New Machine"}
            description="Add custom ComfyUI machines to your account."
            serverAction={addCustomMachine}
            formSchema={addCustomMachineSchema}
@@ -402,11 +434,9 @@ export function MachineList({
                },
              },
              gpu: {
-                fieldType: !userMetadata.betaFeaturesAccess
-                  ? "fallback"
-                  : "select",
+                fieldType: "gpuPicker",
                inputProps: {
-                  disabled: !userMetadata.betaFeaturesAccess,
+                  sub: sub,
                },
              },
            }}
@@ -36,7 +36,7 @@ type State = {
    json: {
      event: string;
      data: any;
-    }
+    },
  ) => void;
 };

@@ -82,7 +82,7 @@ function MachineWS({
  const logs = useStore((x) =>
    x.logs
      .filter((p) => p.machine_id === machine.id)
-      .sort((a, b) => a.timestamp - b.timestamp)
+      .sort((a, b) => a.timestamp - b.timestamp),
  );
  const [sid, setSid] = useState("");

@@ -96,7 +96,7 @@ function MachineWS({
      // queryParams: {
      //   clientId: sid,
      // },
-    }
+    },
  );

  const connectionStatus = getConnectionStatus(readyState);
@@ -135,7 +135,9 @@ function MachineWS({
            You can view your run&apos;s outputs here
          </DialogDescription>
        </DialogHeader>
-        <LogsViewer logs={logs} />
+        <div className="h-[400px]">
+          <LogsViewer logs={logs} hideTimestamp />
+        </div>
      </DialogContent>
    </Dialog>
  );
@@ -7,6 +7,13 @@ import { Checkbox } from "@/components/ui/checkbox";
 import { InsertModal } from "./InsertModal";
 import { Input } from "@/components/ui/input";
 import { ScrollArea } from "@/components/ui/scroll-area";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuLabel,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
 import {
  Table,
  TableBody,
@@ -15,7 +22,7 @@ import {
  TableHeader,
  TableRow,
 } from "@/components/ui/table";
-import type { getAllUserModels as getAllUserModels } from "@/server/getAllUserModel";
+import type { getAllUserModels } from "@/server/getAllUserModel";
 import type {
  ColumnDef,
  ColumnFiltersState,
@@ -30,10 +37,10 @@ import {
  getSortedRowModel,
  useReactTable,
 } from "@tanstack/react-table";
-import { ArrowUpDown } from "lucide-react";
+import { ArrowUpDown, MoreHorizontal } from "lucide-react";
 import * as React from "react";
-import { addCivitaiModel } from "@/server/curdModel";
-import { addCivitaiModelSchema } from "@/server/addCivitaiModelSchema";
+import { addModel, deleteModel } from "@/server/curdModel";
+import { downloadUrlModelSchema } from "@/server/addCivitaiModelSchema";
 import { modelEnumType } from "@/db/schema";

 export type ModelItemList = NonNullable<
@@ -46,8 +53,10 @@ export const columns: ColumnDef<ModelItemList>[] = [
    id: "select",
    header: ({ table }) => (
      <Checkbox
-        checked={table.getIsAllPageRowsSelected() ||
-          (table.getIsSomePageRowsSelected() && "indeterminate")}
+        checked={
+          table.getIsAllPageRowsSelected() ||
+          (table.getIsSomePageRowsSelected() && "indeterminate")
+        }
        onCheckedChange={(value) => table.toggleAllPageRowsSelected(!!value)}
        aria-label="Select all"
      />
@@ -79,19 +88,15 @@ export const columns: ColumnDef<ModelItemList>[] = [
      const model = row.original;
      return (
        <>
-          {
-            /*<a
+          {/*<a
          className="hover:underline flex gap-2"
          href={`/storage/${model.id}`} // TODO
-        >*/
-          }
+        >*/}
          <span className="truncate max-w-[200px]">
            {row.original.model_name}
          </span>

-          {model.is_public
-            ? <></>
-            : <Badge variant="orange">Private</Badge>}
+          {model.is_public ? <></> : <Badge variant="orange">Private</Badge>}
        </>
      );
    },
@@ -112,9 +117,13 @@ export const columns: ColumnDef<ModelItemList>[] = [
    cell: ({ row }) => {
      return (
        <Badge
-          variant={row.original.status === "failed"
-            ? "red"
-            : (row.original.status === "started" ? "yellow" : "green")}
+          variant={
+            row.original.status === "failed"
+              ? "red"
+              : row.original.status === "started"
+                ? "yellow"
+                : "green"
+          }
        >
          {row.original.status}
        </Badge>
@@ -186,14 +195,20 @@ export const columns: ColumnDef<ModelItemList>[] = [
    },
    cell: ({ row }) => {
      const model_type_map: Record<modelEnumType, any> = {
-        "checkpoint": "amber",
-        "lora": "green",
-        "embedding": "violet",
-        "vae": "teal",
+        checkpoint: "amber",
+        lora: "green",
+        embedding: "violet",
+        vae: "teal",
+        clip: "default",
+        clip_vision: "default",
+        configs: "default",
+        controlnet: "default",
+        upscale_models: "default",
+        ipadapter: "default",
      };

      function getBadgeColor(modelType: modelEnumType) {
-        return model_type_map[modelType] || "default";
+        return model_type_map[modelType]
      }

      const color = getBadgeColor(row.original.model_type);
@@ -223,35 +238,35 @@ export const columns: ColumnDef<ModelItemList>[] = [
    ),
  },
  // TODO: deletion and editing for future sprint
-  // {
-  //   id: "actions",
-  //   enableHiding: false,
-  //   cell: ({ row }) => {
-  //     const checkpoint = row.original;
-  //
-  //     return (
-  //       <DropdownMenu>
-  //         <DropdownMenuTrigger asChild>
-  //           <Button variant="ghost" className="h-8 w-8 p-0">
-  //             <span className="sr-only">Open menu</span>
-  //             <MoreHorizontal className="h-4 w-4" />
-  //           </Button>
-  //         </DropdownMenuTrigger>
-  //         <DropdownMenuContent align="end">
-  //           <DropdownMenuLabel>Actions</DropdownMenuLabel>
-  //           <DropdownMenuItem
-  //             className="text-destructive"
-  //             onClick={() => {
-  //               deleteWorkflow(checkpoint.id);
-  //             }}
-  //           >
-  //             Delete Workflow
-  //           </DropdownMenuItem>
-  //         </DropdownMenuContent>
-  //       </DropdownMenu>
-  //     );
-  //   },
-  // },
+  {
+    id: "actions",
+    enableHiding: false,
+    cell: ({ row }) => {
+      const model = row.original;
+
+      return (
+        <DropdownMenu>
+          <DropdownMenuTrigger asChild>
+            <Button variant="ghost" className="h-8 w-8 p-0">
+              <span className="sr-only">Open menu</span>
+              <MoreHorizontal className="h-4 w-4" />
+            </Button>
+          </DropdownMenuTrigger>
+          <DropdownMenuContent align="end">
+            <DropdownMenuLabel>Actions</DropdownMenuLabel>
+            <DropdownMenuItem
+              className="text-destructive"
+              onClick={() => {
+                deleteModel(model.id);
+              }}
+            >
+              Delete Model
+            </DropdownMenuItem>
+          </DropdownMenuContent>
+        </DropdownMenu>
+      );
+    },
+  },
 ];

 export function ModelList({ data }: { data: ModelItemList[] }) {
@@ -259,9 +274,8 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
  const [columnFilters, setColumnFilters] = React.useState<ColumnFiltersState>(
    [],
  );
-  const [columnVisibility, setColumnVisibility] = React.useState<
-    VisibilityState
-  >({});
+  const [columnVisibility, setColumnVisibility] =
+    React.useState<VisibilityState>({});
  const [rowSelection, setRowSelection] = React.useState({});

  const table = useReactTable({
@@ -288,27 +302,27 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
      <div className="flex flex-row w-full items-center py-4">
        <Input
          placeholder="Filter workflows..."
-          value={(table.getColumn("model_name")?.getFilterValue() as string) ??
-            ""}
+          value={
+            (table.getColumn("model_name")?.getFilterValue() as string) ?? ""
+          }
          onChange={(event) =>
-            table.getColumn("model_name")?.setFilterValue(event.target.value)}
+            table.getColumn("model_name")?.setFilterValue(event.target.value)
+          }
          className="max-w-sm"
        />
        <div className="ml-auto flex gap-2">
          <InsertModal
            dialogClassName="sm:max-w-[600px]"
            disabled={
-              false
-              // TODO: limitations based on plan
+              false // TODO: limitations based on plan
            }
-            tooltip={"Add models using their civitai url!"}
-            title="Add a Civitai Model"
-            description="Pick a model from civitai"
-            serverAction={addCivitaiModel}
-            formSchema={addCivitaiModelSchema}
+            title="Add a Model"
+            description="using a link to a model"
+            serverAction={addModel}
+            formSchema={downloadUrlModelSchema}
            fieldConfig={{
-              civitai_url: {
-                fieldType: "fallback",
+              url: {
+                fieldType: "modelUrlPicker",
                inputProps: { required: true },
                description: (
                  <>
@@ -317,13 +331,21 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
                      href="https://www.civitai.com/models"
                      target="_blank"
                      className="underline text-blue-600 hover:text-blue-800 visited:text-purple-600"
+                      rel="noreferrer"
                    >
                      civitai.com
                    </a>{" "}
-                    and place it's url here
+                    or a url we can download a model from
                  </>
                ),
              },
+              model_type: {
+                fieldType: "select",
+                inputProps: { required: true },
+                description: (
+                  <>We'll figure this out if you pick a civitai model</>
+                ),
+              },
            }}
          />
        </div>
@@ -336,10 +358,12 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
                {headerGroup.headers.map((header) => {
                  return (
                    <TableHead key={header.id}>
-                      {header.isPlaceholder ? null : flexRender(
-                        header.column.columnDef.header,
-                        header.getContext(),
-                      )}
+                      {header.isPlaceholder
+                        ? null
+                        : flexRender(
+                            header.column.columnDef.header,
+                            header.getContext(),
+                          )}
                    </TableHead>
                  );
                })}
@@ -347,34 +371,32 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
            ))}
          </TableHeader>
          <TableBody>
-            {table.getRowModel().rows?.length
-              ? (
-                table.getRowModel().rows.map((row) => (
-                  <TableRow
-                    key={row.id}
-                    data-state={row.getIsSelected() && "selected"}
-                  >
-                    {row.getVisibleCells().map((cell) => (
-                      <TableCell key={cell.id}>
-                        {flexRender(
-                          cell.column.columnDef.cell,
-                          cell.getContext(),
-                        )}
-                      </TableCell>
-                    ))}
-                  </TableRow>
-                ))
-              )
-              : (
-                <TableRow>
-                  <TableCell
-                    colSpan={columns.length}
-                    className="h-24 text-center"
-                  >
-                    No results.
-                  </TableCell>
+            {table.getRowModel().rows?.length ? (
+              table.getRowModel().rows.map((row) => (
+                <TableRow
+                  key={row.id}
+                  data-state={row.getIsSelected() && "selected"}
+                >
+                  {row.getVisibleCells().map((cell) => (
+                    <TableCell key={cell.id}>
+                      {flexRender(
+                        cell.column.columnDef.cell,
+                        cell.getContext(),
+                      )}
+                    </TableCell>
+                  ))}
                </TableRow>
-              )}
+              ))
+            ) : (
+              <TableRow>
+                <TableCell
+                  colSpan={columns.length}
+                  className="h-24 text-center"
+                >
+                  No results.
+                </TableCell>
+              </TableRow>
+            )}
          </TableBody>
        </Table>
      </ScrollArea>
@@ -1,27 +1,42 @@
+"use client";
+
+import useSWR from "swr";
 import { DownloadButton } from "./DownloadButton";
 import { getFileDownloadUrl } from "@/server/getFileDownloadUrl";

-export async function OutputRender(props: {
+export function OutputRender(props: {
  run_id: string;
  filename: string;
 }) {
+  const { data: url } = useSWR(
+    "run-outputs+" + props.run_id + props.filename,
+    async () => {
+      return await getFileDownloadUrl(
+        `outputs/runs/${props.run_id}/${props.filename}`,
+      );
+    },
+  );
+
+  if (!url) return <></>;
+
+  if (props.filename.endsWith(".mp4") || props.filename.endsWith(".webm")) {
+    return (
+      <video controls autoPlay className="w-[400px]">
+        <source src={url} type="video/mp4" />
+        <source src={url} type="video/webm" />
+        Your browser does not support the video tag.
+      </video>
+    );
+  }
+
  if (
    props.filename.endsWith(".png") ||
    props.filename.endsWith(".gif") ||
    props.filename.endsWith(".jpg") ||
    props.filename.endsWith(".jpeg")
  ) {
-    const url = await getFileDownloadUrl(
-      `outputs/runs/${props.run_id}/${props.filename}`
-    );
-
    return <img className="max-w-[200px]" alt={props.filename} src={url} />;
  } else {
-    const url = await getFileDownloadUrl(
-      `outputs/runs/${props.run_id}/${props.filename}`
-    );
-    // console.log(url);
-
    return <DownloadButton filename={props.filename} href={url} />;
  }
 }
@@ -1,3 +1,5 @@
+"use client";
+
 import { RunInputs } from "@/components/RunInputs";
 import { RunOutputs } from "@/components/RunOutputs";
 import { Badge } from "@/components/ui/badge";
@@ -17,10 +19,10 @@ import {
 } from "@/components/ui/tooltip";
 import { getDuration, getRelativeTime } from "@/lib/getRelativeTime";
 import { type findAllRuns } from "@/server/findAllRuns";
-import { Suspense } from "react";
 import { LiveStatus } from "./LiveStatus";
+import { LoadingWrapper } from "@/components/LoadingWrapper";

-export async function RunDisplay({
+export function RunDisplay({
  run,
 }: {
  run: Awaited<ReturnType<typeof findAllRuns>>[0];
@@ -73,9 +75,9 @@ export async function RunDisplay({
        </DialogHeader>
        <div className="max-h-96 overflow-y-scroll">
          <RunInputs run={run} />
-          <Suspense>
-            <RunOutputs run_id={run.id} />
-          </Suspense>
+          <LoadingWrapper tag="output">
+            <RunOutputs run={run} />
+          </LoadingWrapper>
        </div>
        {/* <div className="max-h-96 overflow-y-scroll">{view}</div> */}
      </DialogContent>
@@ -8,7 +8,7 @@ import {
 } from "@/components/ui/table";
 import type { findAllRuns } from "@/server/findAllRuns";

-export async function RunInputs({
+export function RunInputs({
  run,
 }: {
  run: Awaited<ReturnType<typeof findAllRuns>>[0];
@@ -1,5 +1,16 @@
+"use client";
+
+import {
+  Dialog,
+  DialogClose,
+  DialogContent,
+  DialogDescription,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+  DialogTrigger,
+} from "@/components/ui/dialog";
 import { OutputRender } from "./OutputRender";
-import { CodeBlock } from "@/components/CodeBlock";
 import {
  Table,
  TableBody,
@@ -8,10 +19,25 @@ import {
  TableHeader,
  TableRow,
 } from "@/components/ui/table";
+import type { findAllRuns } from "@/server/findAllRuns";
 import { getRunsOutput } from "@/server/getRunsOutput";
+import { Button } from "@/components/ui/button";
+import { ExternalLink } from "lucide-react";
+import { LogsViewer } from "@/components/LogsViewer";
+import { CopyButton } from "@/components/CopyButton";
+import useSWR from "swr";
+import { CodeBlockClient } from "@/components/CodeBlockClient";
+
+export function RunOutputs({
+  run,
+}: { run: Awaited<ReturnType<typeof findAllRuns>>[0] }) {
+  const { data, isValidating, error } = useSWR(
+    "run-outputs+" + run.id,
+    async () => {
+      return await getRunsOutput(run.id);
+    },
+  );

-export async function RunOutputs({ run_id }: { run_id: string }) {
-  const outputs = await getRunsOutput(run_id);
  return (
    <Table className="table-fixed">
      <TableHeader className="bg-background top-0 sticky">
@@ -21,7 +47,43 @@ export async function RunOutputs({ run_id }: { run_id: string }) {
        </TableRow>
      </TableHeader>
      <TableBody>
-        {outputs?.map((run) => {
+        <TableRow key={run.id}>
+          <TableCell className="break-words">Run log</TableCell>
+          <TableCell>
+            {run.run_log ? (
+              <Dialog>
+                <DialogTrigger asChild>
+                  <Button variant="secondary" className="w-fit">
+                    View Log <ExternalLink size={14} />
+                  </Button>
+                </DialogTrigger>
+                <DialogContent className="sm:max-w-[1000px] h-full max-h-[600px] grid grid-rows-[auto,1fr,auto]">
+                  <DialogHeader>
+                    <DialogTitle>Run Log</DialogTitle>
+                  </DialogHeader>
+                  <LogsViewer logs={run.run_log} stickToBottom={false} />
+                  <DialogFooter>
+                    <CopyButton
+                      className="w-fit aspect-auto p-4"
+                      text={JSON.stringify(run.run_log)}
+                    >
+                      Copy
+                    </CopyButton>
+                    <DialogClose>
+                      <Button type="button" variant="secondary">
+                        Close
+                      </Button>
+                    </DialogClose>
+                  </DialogFooter>
+                </DialogContent>
+              </Dialog>
+            ) : (
+              "No log available"
+            )}
+          </TableCell>
+        </TableRow>
+
+        {data?.map((run) => {
          const fileName =
            run.data.images?.[0].filename ||
            run.data.files?.[0].filename ||
@@ -32,7 +94,7 @@ export async function RunOutputs({ run_id }: { run_id: string }) {
              <TableRow key={run.id}>
                <TableCell>Output</TableCell>
                <TableCell className="">
-                  <CodeBlock
+                  <CodeBlockClient
                    code={JSON.stringify(run.data, null, 2)}
                    lang="json"
                  />
@@ -45,7 +107,7 @@ export async function RunOutputs({ run_id }: { run_id: string }) {
            <TableRow key={run.id}>
              <TableCell className="break-words">{fileName}</TableCell>
              <TableCell>
-                <OutputRender run_id={run_id} filename={fileName} />
+                <OutputRender run_id={run.run_id} filename={fileName} />
              </TableCell>
            </TableRow>
          );
@@ -1,3 +1,5 @@
+"use client";
+
 import {
  Table,
  TableBody,
@@ -7,94 +9,79 @@ import {
  TableRow,
 } from "@/components/ui/table";
 import { parseAsInteger } from "next-usequerystate";
-import { headers } from "next/headers";
 import {
-	findAllDeployments,
-	findAllRunsWithCounts,
+  findAllRunsWithCounts,
+  getAllRunstableContent,
 } from "../server/findAllRuns";
-import { DeploymentDisplay } from "./DeploymentDisplay";
 import { PaginationControl } from "./PaginationControl";
 import { RunDisplay } from "./RunDisplay";
+import useSWR from "swr";
+import { LoadingIcon } from "@/components/LoadingIcon";

 const itemPerPage = 6;
 const pageParser = parseAsInteger.withDefault(1);

-export async function RunsTable(props: {
+export function RunsTable(props: {
  workflow_id: string;
-  searchParams: { [key: string]: string | string[] | undefined };
+  searchParams: { [key: string]: any };
 }) {
-  // await new Promise((resolve) => setTimeout(resolve, 5000));
-  const page = pageParser.parseServerSide(
-    props.searchParams?.page ?? undefined
+  const page = pageParser.parse(props.searchParams?.page ?? undefined) ?? 1;
+  const { data, error, isLoading, isValidating } = useSWR(
+    "runs+" + page,
+    async () => {
+      const data = await findAllRunsWithCounts({
+        workflow_id: props.workflow_id,
+        limit: itemPerPage,
+        offset: (page - 1) * itemPerPage,
+      });
+      return data;
+    },
+    {
+      // suspense: false,
+      refreshInterval: 5000,
+    },
  );
-  const { allRuns, total } = await findAllRunsWithCounts({
-    workflow_id: props.workflow_id,
-    limit: itemPerPage,
-    offset: (page - 1) * itemPerPage,
-  });
-  return (
-			<div>
-				<div className="overflow-auto h-fit w-full">
-					<Table className="">
-						{allRuns.length === 0 && (
-							<TableCaption>A list of your recent runs.</TableCaption>
-						)}
-						<TableHeader className="bg-background top-0 sticky">
-							<TableRow>
-								<TableHead className="truncate">Number</TableHead>
-								<TableHead className="truncate">Machine</TableHead>
-								<TableHead className="truncate">Time</TableHead>
-								<TableHead className="truncate">Version</TableHead>
-								<TableHead className="truncate">Origin</TableHead>
-								<TableHead className="truncate">Duration</TableHead>
-								<TableHead className="truncate">Live Status</TableHead>
-								<TableHead className="text-right">Status</TableHead>
-							</TableRow>
-						</TableHeader>
-						<TableBody>
-							{allRuns.map((run) => (
-								<RunDisplay run={run} key={run.id} />
-							))}
-						</TableBody>
-					</Table>
-				</div>

-				{Math.ceil(total / itemPerPage) > 0 && (
-					<PaginationControl
-						totalPage={Math.ceil(total / itemPerPage)}
-						currentPage={page}
-					/>
-				)}
-			</div>
-		);
-}
-
-export async function DeploymentsTable(props: { workflow_id: string }) {
-  const allRuns = await findAllDeployments(props.workflow_id);
-
-  const headersList = headers();
-  const host = headersList.get("host") || "";
-  const protocol = headersList.get("x-forwarded-proto") || "";
-  const domain = `${protocol}://${host}`;
+  // await new Promise((resolve) => setTimeout(resolve, 5000));

  return (
-    <div className="overflow-auto h-fit w-full">
-      <Table className="">
-        <TableCaption>A list of your deployments</TableCaption>
-        <TableHeader className="bg-background top-0 sticky">
-          <TableRow>
-            <TableHead className=" w-[100px]">Environment</TableHead>
-            <TableHead className=" w-[100px]">Version</TableHead>
-            <TableHead className="">Machine</TableHead>
-            <TableHead className=" text-right">Updated At</TableHead>
-          </TableRow>
-        </TableHeader>
-        <TableBody>
-          {allRuns.map((run) => (
-            <DeploymentDisplay deployment={run} key={run.id} domain={domain} />
-          ))}
-        </TableBody>
-      </Table>
+    <div>
+      {isValidating ? (
+        <div className="absolute right-8 top-8">
+          <LoadingIcon />
+        </div>
+      ) : null}
+      <div className="overflow-auto h-fit w-full">
+        <Table className="">
+          {/* {data?.allRuns.length === 0 && (
+            <TableCaption>A list of your recent runs.</TableCaption>
+          )} */}
+          <TableHeader className="bg-background top-0 sticky">
+            <TableRow>
+              <TableHead className="truncate">Number</TableHead>
+              <TableHead className="truncate">Machine</TableHead>
+              <TableHead className="truncate">Time</TableHead>
+              <TableHead className="truncate">Version</TableHead>
+              <TableHead className="truncate">Origin</TableHead>
+              <TableHead className="truncate">Duration</TableHead>
+              <TableHead className="truncate">Live Status</TableHead>
+              <TableHead className="text-right">Status</TableHead>
+            </TableRow>
+          </TableHeader>
+          <TableBody>
+            {data?.allRuns.map((run) => (
+              <RunDisplay run={run} key={run.id} />
+            ))}
+          </TableBody>
+        </Table>
+      </div>
+
+      {data && Math.ceil(data.total / itemPerPage) > 0 && (
+        <PaginationControl
+          totalPage={Math.ceil(data.total / itemPerPage)}
+          currentPage={page}
+        />
+      )}
    </div>
  );
 }
@@ -17,6 +17,8 @@ export function StatusBadge({
      );
    case "success":
      return <Badge variant="success">{status}</Badge>;
+    case "timeout":
+      return <Badge variant="amber">{status}</Badge>;
    case "failed":
      return <Badge variant="destructive">{status}</Badge>;
  }
@@ -2,7 +2,6 @@

 import { LoadingIcon } from "@/components/LoadingIcon";
 import AutoForm, { AutoFormSubmit } from "@/components/ui/auto-form";
-import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
 import {
  Dialog,
@@ -28,38 +27,22 @@ import {
  SelectValue,
 } from "@/components/ui/select";
 import { Skeleton } from "@/components/ui/skeleton";
-import {
-  Table,
-  TableBody,
-  TableCell,
-  TableHead,
-  TableHeader,
-  TableRow,
-} from "@/components/ui/table";
-import type { showcaseMediaNullable, workflowAPINodeType } from "@/db/schema";
+import type { showcaseMediaNullable } from "@/db/schema";
 import { checkStatus, createRun } from "@/server/createRun";
 import { createDeployments } from "@/server/curdDeploments";
 import type { getMachines } from "@/server/curdMachine";
 import type { findFirstTableWithVersion } from "@/server/findFirstTableWithVersion";
-import {
-  Copy,
-  Edit,
-  ExternalLink,
-  Info,
-  MoreVertical,
-  Play,
-} from "lucide-react";
+import { Copy, Edit, MoreVertical, Play } from "lucide-react";
 import { parseAsInteger, useQueryState } from "next-usequerystate";
-import { useEffect, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useState } from "react";
 import { toast } from "sonner";
-import useSWR from "swr";
 import type { z } from "zod";
 import { create } from "zustand";
 import { workflowVersionInputsToZod } from "../lib/workflowVersionInputsToZod";
 import { callServerPromise } from "./callServerPromise";
-import fetcher from "./fetcher";
 import { ButtonAction } from "@/components/ButtonActionLoader";
 import { editWorkflowOnMachine } from "@/server/editWorkflowOnMachine";
+import { usePathname, useRouter, useSearchParams } from "next/navigation";

 export function VersionSelect({
  workflow,
@@ -124,14 +107,48 @@ export function MachineSelect({
  );
 }

+type SelectedMachineStore = {
+  selectedMachine: string | undefined;
+  setSelectedMachine: (machine: string) => void;
+};
+
+export const selectedMachineStore = create<SelectedMachineStore>((set) => ({
+  selectedMachine: undefined,
+  setSelectedMachine: (machine) => set(() => ({ selectedMachine: machine })),
+}));
+
 export function useSelectedMachine(
  machines: Awaited<ReturnType<typeof getMachines>>,
-) {
-  const a = useQueryState("machine", {
-    defaultValue: machines?.[0]?.id ?? "",
-  });
+): [string, (v: string) => void] {
+  const { selectedMachine, setSelectedMachine } = selectedMachineStore();
+  return [selectedMachine ?? machines?.[0]?.id ?? "", setSelectedMachine];

-  return a;
+  // const searchParams = useSearchParams();
+  // const pathname = usePathname();
+  // const router = useRouter();
+
+  // const createQueryString = useCallback(
+  //   (name: string, value: string) => {
+  //     const params = new URLSearchParams(searchParams.toString());
+  //     params.set(name, value);
+
+  //     return params.toString();
+  //   },
+  //   [searchParams],
+  // );
+
+  // return [
+  //   searchParams.get("machine") ?? machines?.[0]?.id ?? "",
+  //   (v: string) => {
+  //     // window.history.pushState(
+  //     //   "new url",
+  //     //   "",
+  //     //   pathname + "?" + createQueryString("machine", v),
+  //     // );
+  //     // router.push(pathname + "?" + createQueryString("machine", v));
+  //     router.replace(pathname + "?" + createQueryString("machine", v));
+  //   },
+  // ];
 }

 type PublicRunStore = {
@@ -484,150 +501,3 @@ export function getWorkflowVersionFromVersionIndex(

  return workflow_version;
 }
-
-export function ViewWorkflowDetailsButton({
-  workflow,
-}: {
-  workflow: Awaited<ReturnType<typeof findFirstTableWithVersion>>;
-}) {
-  const [version] = useQueryState("version", {
-    defaultValue: workflow?.versions[0].version ?? 1,
-    ...parseAsInteger,
-  });
-  const [isLoading, setIsLoading] = useState(false);
-
-  const [open, setOpen] = useState(false);
-
-  const {
-    data,
-    error,
-    isLoading: isNodesIndexLoading,
-  } = useSWR(
-    "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/extension-node-map.json",
-    fetcher,
-  );
-
-  const groupedByAuxName = useMemo(() => {
-    if (!data) return null;
-
-    // console.log(data);
-
-    const workflow_version = getWorkflowVersionFromVersionIndex(
-      workflow,
-      version,
-    );
-
-    const api = workflow_version?.workflow_api;
-
-    if (!api) return null;
-
-    const crossCheckedApi = Object.entries(api)
-      .map(([_, value]) => {
-        const classType = value.class_type;
-        const classTypeData = Object.entries(data).find(([_, nodeArray]) =>
-          nodeArray[0].includes(classType),
-        );
-        return classTypeData ? { node: value, classTypeData } : null;
-      })
-      .filter((item) => item !== null);
-
-    // console.log(crossCheckedApi);
-
-    const groupedByAuxName = crossCheckedApi.reduce(
-      (acc, data) => {
-        if (!data) return acc;
-
-        const { node, classTypeData } = data;
-        const auxName = classTypeData[1][1].title_aux;
-        // console.log(auxName);
-        if (!acc[auxName]) {
-          acc[auxName] = {
-            url: classTypeData[0],
-            node: [],
-          };
-        }
-        acc[auxName].node.push(node);
-        return acc;
-      },
-      {} as Record<
-        string,
-        {
-          node: z.infer<typeof workflowAPINodeType>[];
-          url: string;
-        }
-      >,
-    );
-
-    // console.log(groupedByAuxName);
-
-    return groupedByAuxName;
-  }, [version, data]);
-
-  return (
-    <Dialog open={open} onOpenChange={setOpen}>
-      <DialogTrigger asChild className="appearance-none hover:cursor-pointer">
-        <Button className="gap-2" variant="outline">
-          Details <Info size={14} />
-        </Button>
-      </DialogTrigger>
-      <DialogContent className="max-w-xl">
-        <DialogHeader>
-          <DialogTitle>Workflow Details</DialogTitle>
-          <DialogDescription>
-            View your custom nodes, models, external files used in this workflow
-          </DialogDescription>
-        </DialogHeader>
-
-        <div className="overflow-auto max-h-[400px] w-full">
-          <Table>
-            <TableHeader className="bg-background top-0 sticky">
-              <TableRow>
-                <TableHead className="w-[200px]">File</TableHead>
-                <TableHead className="">Output</TableHead>
-              </TableRow>
-            </TableHeader>
-            <TableBody>
-              {groupedByAuxName &&
-                Object.entries(groupedByAuxName).map(([key, group]) => {
-                  // const filePath
-                  return (
-                    <TableRow key={key}>
-                      <TableCell className="break-words">
-                        <a
-                          href={group.url}
-                          target="_blank"
-                          className="hover:underline"
-                          rel="noreferrer"
-                        >
-                          {key}
-                          <ExternalLink
-                            className="inline-block ml-1"
-                            size={12}
-                          />
-                        </a>
-                      </TableCell>
-                      <TableCell className="flex flex-wrap gap-2">
-                        {group.node.map((x) => (
-                          <Badge key={x.class_type} variant="outline">
-                            {x.class_type}
-                          </Badge>
-                        ))}
-                      </TableCell>
-                    </TableRow>
-                  );
-                })}
-            </TableBody>
-          </Table>
-        </div>
-
-        <div className="flex justify-end">
-          <Button className="w-fit" onClick={() => setOpen(false)}>
-            Close
-          </Button>
-        </div>
-        {/* </div> */}
-        {/* <div className="max-h-96 overflow-y-scroll">{view}</div> */}
-      </DialogContent>
-    </Dialog>
-  );
-}
@@ -0,0 +1,169 @@
+"use client";
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import {
+  Dialog,
+  DialogClose,
+  DialogContent,
+  DialogDescription,
+  DialogHeader,
+  DialogTitle,
+  DialogTrigger,
+} from "@/components/ui/dialog";
+import {
+  Table,
+  TableBody,
+  TableCell,
+  TableHead,
+  TableHeader,
+  TableRow,
+} from "@/components/ui/table";
+import type { workflowAPINodeType } from "@/db/schema";
+import type { findFirstTableWithVersion } from "@/server/findFirstTableWithVersion";
+import { ExternalLink, Info } from "lucide-react";
+import { parseAsInteger, useQueryState } from "next-usequerystate";
+import { useMemo } from "react";
+import useSWR from "swr";
+import type { z } from "zod";
+import fetcher from "./fetcher";
+import { getWorkflowVersionFromVersionIndex } from "./VersionSelect";
+
+export function ViewWorkflowDetailsButton({
+  workflow,
+}: {
+  workflow: Awaited<ReturnType<typeof findFirstTableWithVersion>>;
+}) {
+  const [version] = useQueryState("version", {
+    defaultValue: workflow?.versions[0].version ?? 1,
+    ...parseAsInteger,
+  });
+
+  const {
+    data,
+    error,
+    isLoading: isNodesIndexLoading,
+  } = useSWR(
+    "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/extension-node-map.json",
+    fetcher,
+  );
+
+  const groupedByAuxName = useMemo(() => {
+    if (!data) return null;
+
+    const workflow_version = getWorkflowVersionFromVersionIndex(
+      workflow,
+      version,
+    );
+
+    const api = workflow_version?.workflow_api;
+
+    if (!api) return null;
+
+    const crossCheckedApi = Object.entries(api)
+      .map(([_, value]) => {
+        const classType = value.class_type;
+        const classTypeData = Object.entries(data).find(([_, nodeArray]) =>
+          nodeArray[0].includes(classType),
+        );
+        return classTypeData ? { node: value, classTypeData } : null;
+      })
+      .filter((item) => item !== null);
+
+    // console.log(crossCheckedApi);
+    const groupedByAuxName = crossCheckedApi.reduce(
+      (acc, data) => {
+        if (!data) return acc;
+
+        const { node, classTypeData } = data;
+        const auxName = classTypeData[1][1].title_aux;
+        // console.log(auxName);
+        if (!acc[auxName]) {
+          acc[auxName] = {
+            url: classTypeData[0],
+            node: [],
+          };
+        }
+        acc[auxName].node.push(node);
+        return acc;
+      },
+      {} as Record<
+        string,
+        {
+          node: z.infer<typeof workflowAPINodeType>[];
+          url: string;
+        }
+      >,
+    );
+
+    // console.log(groupedByAuxName);
+    return groupedByAuxName;
+  }, [version, data]);
+
+  return (
+    <Dialog>
+      <DialogTrigger asChild className="appearance-none hover:cursor-pointer">
+        <Button className="gap-2" variant="outline">
+          Details <Info size={14} />
+        </Button>
+      </DialogTrigger>
+      <DialogContent className="max-w-xl">
+        <DialogHeader>
+          <DialogTitle>Workflow Details</DialogTitle>
+          <DialogDescription>
+            View your custom nodes, models, external files used in this workflow
+          </DialogDescription>
+        </DialogHeader>
+
+        <div className="overflow-auto max-h-[400px] w-full">
+          <Table>
+            <TableHeader className="bg-background top-0 sticky">
+              <TableRow>
+                <TableHead className="w-[200px]">File</TableHead>
+                <TableHead className="">Output</TableHead>
+              </TableRow>
+            </TableHeader>
+            <TableBody>
+              {groupedByAuxName &&
+                Object.entries(groupedByAuxName).map(([key, group]) => {
+                  // const filePath
+                  return (
+                    <TableRow key={key}>
+                      <TableCell className="break-words">
+                        <a
+                          href={group.url}
+                          target="_blank"
+                          className="hover:underline"
+                          rel="noreferrer"
+                        >
+                          {key}
+                          <ExternalLink
+                            className="inline-block ml-1"
+                            size={12}
+                          />
+                        </a>
+                      </TableCell>
+                      <TableCell className="flex flex-wrap gap-2">
+                        {group.node.map((x) => (
+                          <Badge key={x.class_type} variant="outline">
+                            {x.class_type}
+                          </Badge>
+                        ))}
+                      </TableCell>
+                    </TableRow>
+                  );
+                })}
+            </TableBody>
+          </Table>
+        </div>
+
+        <div className="flex justify-end">
+          <DialogClose asChild>
+            <Button className="w-fit">Close</Button>
+          </DialogClose>
+        </div>
+        {/* </div> */}
+        {/* <div className="max-h-96 overflow-y-scroll">{view}</div> */}
+      </DialogContent>
+    </Dialog>
+  );
+}
@@ -0,0 +1,72 @@
+"use client";
+import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
+import * as React from "react";
+import { useDebouncedCallback } from "use-debounce";
+import { z } from "zod";
+import { CivitalModelSchema, ModelListWrapper } from "./CivitalModelSchema";
+import { getUrl, mapModelsList } from "./getUrl";
+import { ModelSelector } from "./ModelSelector";
+
+export function CivitaiModelRegistry({
+  field,
+  selectMultiple = true,
+}: Pick<AutoFormInputComponentProps, "field"> & {
+  selectMultiple?: boolean;
+}) {
+  const [modelList, setModelList] =
+    React.useState<z.infer<typeof ModelListWrapper>>();
+
+  const [loading, setLoading] = React.useState(false);
+
+  const handleSearch = useDebouncedCallback((search) => {
+    console.log(`Searching... ${search}`);
+
+    setLoading(true);
+
+    const controller = new AbortController();
+    fetch(getUrl(search), {
+      signal: controller.signal,
+    })
+      .then((x) => x.json())
+      .then((a) => {
+        const list = CivitalModelSchema.parse(a);
+        console.log(a);
+
+        setModelList(mapModelsList(list));
+        setLoading(false);
+      });
+
+    return () => {
+      controller.abort();
+      setLoading(false);
+    };
+  }, 300);
+
+  React.useEffect(() => {
+    const controller = new AbortController();
+    fetch(getUrl(), {
+      signal: controller.signal,
+    })
+      .then((x) => x.json())
+      .then((a) => {
+        const list = CivitalModelSchema.parse(a);
+        setModelList(mapModelsList(list));
+      });
+
+    return () => {
+      controller.abort();
+    };
+  }, []);
+
+  return (
+    <ModelSelector
+      selectMultiple={selectMultiple}
+      field={field}
+      modelList={modelList}
+      label="Civitai"
+      onSearch={handleSearch}
+      shouldFilter={false}
+      isLoading={loading}
+    />
+  );
+}
@@ -0,0 +1,94 @@
+"use client";
+import { z } from "zod";
+
+export const Model = z.object({
+  name: z.string(),
+  type: z.string(),
+  base: z.string(),
+  save_path: z.string(),
+  description: z.string(),
+  reference: z.string(),
+  filename: z.string(),
+  url: z.string(),
+});
+
+export const CivitalModelSchema = z.object({
+  items: z.array(
+    z.object({
+      id: z.number(),
+      name: z.string(),
+      description: z.string(),
+      type: z.string(),
+      creator: z
+        .object({
+          username: z.string().nullable(),
+          image: z.string().nullable().default(null),
+        })
+        .nullable(),
+      tags: z.array(z.string()),
+      modelVersions: z.array(
+        z.object({
+          id: z.number(),
+          modelId: z.number(),
+          name: z.string(),
+          createdAt: z.string(),
+          updatedAt: z.string(),
+          status: z.string(),
+          publishedAt: z.string(),
+          trainedWords: z.array(z.unknown()),
+          trainingStatus: z.string().nullable(),
+          trainingDetails: z.string().nullable(),
+          baseModel: z.string(),
+          baseModelType: z.string().nullable(),
+          earlyAccessTimeFrame: z.number(),
+          description: z.string().nullable(),
+          vaeId: z.number().nullable(),
+          stats: z.object({
+            downloadCount: z.number(),
+            ratingCount: z.number(),
+            rating: z.number(),
+          }),
+          files: z.array(
+            z.object({
+              id: z.number(),
+              sizeKB: z.number(),
+              name: z.string(),
+              type: z.string(),
+              downloadUrl: z.string(),
+            }),
+          ),
+          images: z.array(
+            z.object({
+              id: z.number(),
+              url: z.string(),
+              nsfw: z.string(),
+              width: z.number(),
+              height: z.number(),
+              hash: z.string(),
+              type: z.string(),
+              metadata: z.object({
+                hash: z.string(),
+                width: z.number(),
+                height: z.number(),
+              }),
+              meta: z.any(),
+            }),
+          ),
+          downloadUrl: z.string(),
+        }),
+      ),
+    }),
+  ),
+  metadata: z.object({
+    totalItems: z.number(),
+    currentPage: z.number(),
+    pageSize: z.number(),
+    totalPages: z.number(),
+    nextPage: z.string().optional(),
+  }),
+});
+export const ModelList = z.array(Model);
+
+export const ModelListWrapper = z.object({
+  models: ModelList,
+});
@@ -0,0 +1,43 @@
+"use client";
+import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
+import * as React from "react";
+import { z } from "zod";
+import { ModelListWrapper } from "./CivitalModelSchema";
+import { ModelSelector } from "./ModelSelector";
+
+export function ComfyUIManagerModelRegistry({
+  field,
+  selectMultiple = true,
+}: Pick<AutoFormInputComponentProps, "field"> & {
+  selectMultiple?: boolean;
+}) {
+  const [modelList, setModelList] =
+    React.useState<z.infer<typeof ModelListWrapper>>();
+
+  React.useEffect(() => {
+    const controller = new AbortController();
+    fetch(
+      "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/model-list.json",
+      {
+        signal: controller.signal,
+      },
+    )
+      .then((x) => x.json())
+      .then((a) => {
+        setModelList(ModelListWrapper.parse(a));
+      });
+
+    return () => {
+      controller.abort();
+    };
+  }, []);
+
+  return (
+    <ModelSelector
+      selectMultiple={selectMultiple}
+      field={field}
+      modelList={modelList}
+      label="ComfyUI Manager"
+    />
+  );
+}
@@ -1,160 +1,17 @@
 "use client";

 import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
-import { LoadingIcon } from "@/components/LoadingIcon";
 import {
  Accordion,
  AccordionContent,
  AccordionItem,
  AccordionTrigger,
 } from "@/components/ui/accordion";
-import { Button } from "@/components/ui/button";
-import {
-  Command,
-  CommandEmpty,
-  CommandGroup,
-  CommandInput,
-  CommandItem,
-  CommandList,
-} from "@/components/ui/command";
-import {
-  Popover,
-  PopoverContent,
-  PopoverTrigger,
-} from "@/components/ui/popover";
 import { ScrollArea } from "@/components/ui/scroll-area";
 import { Textarea } from "@/components/ui/textarea";
-import { cn } from "@/lib/utils";
-import { Check, ChevronsUpDown } from "lucide-react";
 import * as React from "react";
-import { useRef } from "react";
-import { useDebouncedCallback } from "use-debounce";
-import { z } from "zod";
-
-const Model = z.object({
-  name: z.string(),
-  type: z.string(),
-  base: z.string(),
-  save_path: z.string(),
-  description: z.string(),
-  reference: z.string(),
-  filename: z.string(),
-  url: z.string(),
-});
-
-export const CivitalModelSchema = z.object({
-  items: z.array(
-    z.object({
-      id: z.number(),
-      name: z.string(),
-      description: z.string(),
-      type: z.string(),
-      // poi: z.boolean(),
-      // nsfw: z.boolean(),
-      // allowNoCredit: z.boolean(),
-      // allowCommercialUse: z.string(),
-      // allowDerivatives: z.boolean(),
-      // allowDifferentLicense: z.boolean(),
-      // stats: z.object({
-      //   downloadCount: z.number(),
-      //   favoriteCount: z.number(),
-      //   commentCount: z.number(),
-      //   ratingCount: z.number(),
-      //   rating: z.number(),
-      //   tippedAmountCount: z.number(),
-      // }),
-      creator: z
-        .object({
-          username: z.string().nullable(),
-          image: z.string().nullable().default(null),
-        })
-        .nullable(),
-      tags: z.array(z.string()),
-      modelVersions: z.array(
-        z.object({
-          id: z.number(),
-          modelId: z.number(),
-          name: z.string(),
-          createdAt: z.string(),
-          updatedAt: z.string(),
-          status: z.string(),
-          publishedAt: z.string(),
-          trainedWords: z.array(z.unknown()),
-          trainingStatus: z.string().nullable(),
-          trainingDetails: z.string().nullable(),
-          baseModel: z.string(),
-          baseModelType: z.string().nullable(),
-          earlyAccessTimeFrame: z.number(),
-          description: z.string().nullable(),
-          vaeId: z.number().nullable(),
-          stats: z.object({
-            downloadCount: z.number(),
-            ratingCount: z.number(),
-            rating: z.number(),
-          }),
-          files: z.array(
-            z.object({
-              id: z.number(),
-              sizeKB: z.number(),
-              name: z.string(),
-              type: z.string(),
-              // metadata: z.object({
-              //   fp: z.string().nullable().optional(),
-              //   size: z.string().nullable().optional(),
-              //   format: z.string().nullable().optional(),
-              // }),
-              // pickleScanResult: z.string(),
-              // pickleScanMessage: z.string(),
-              // virusScanResult: z.string(),
-              // virusScanMessage: z.string().nullable(),
-              // scannedAt: z.string(),
-              // hashes: z.object({
-              //   AutoV1: z.string().nullable().optional(),
-              //   AutoV2: z.string().nullable().optional(),
-              //   SHA256: z.string().nullable().optional(),
-              //   CRC32: z.string().nullable().optional(),
-              //   BLAKE3: z.string().nullable().optional(),
-              // }),
-              downloadUrl: z.string(),
-              // primary: z.boolean().default(false),
-            })
-          ),
-          images: z.array(
-            z.object({
-              id: z.number(),
-              url: z.string(),
-              nsfw: z.string(),
-              width: z.number(),
-              height: z.number(),
-              hash: z.string(),
-              type: z.string(),
-              metadata: z.object({
-                hash: z.string(),
-                width: z.number(),
-                height: z.number(),
-              }),
-              meta: z.any(),
-            })
-          ),
-          downloadUrl: z.string(),
-        })
-      ),
-    })
-  ),
-  metadata: z.object({
-    totalItems: z.number(),
-    currentPage: z.number(),
-    pageSize: z.number(),
-    totalPages: z.number(),
-    nextPage: z.string().optional(),
-  }),
-});
-
-const ModelList = z.array(Model);
-
-export const ModelListWrapper = z.object({
-  models: ModelList,
-});
+import { CivitaiModelRegistry } from "./CivitaiModelRegistry";
+import { ComfyUIManagerModelRegistry } from "./ComfyUIManagerModelRegistry";

 export function ModelPickerView({
  field,
@@ -187,240 +44,3 @@ export function ModelPickerView({
    </Accordion>
  );
 }
-
-function mapType(type: string) {
-  switch (type) {
-    case "checkpoint":
-      return "checkpoints";
-  }
-  return type;
-}
-
-function mapModelsList(
-  models: z.infer<typeof CivitalModelSchema>
-): z.infer<typeof ModelListWrapper> {
-  return {
-    models: models.items.flatMap((item) => {
-      return item.modelVersions.map((v) => {
-        return {
-          name: `${item.name} ${v.name} (${v.files[0].name})`,
-          type: mapType(item.type.toLowerCase()),
-          base: v.baseModel,
-          save_path: "default",
-          description: item.description,
-          reference: "",
-          filename: v.files[0].name,
-          url: v.files[0].downloadUrl,
-        } as z.infer<typeof Model>;
-      });
-    }),
-  };
-}
-
-function getUrl(search?: string) {
-  const baseUrl = "https://civitai.com/api/v1/models";
-  const searchParams = {
-    limit: 5,
-  } as any;
-  searchParams["sort"] = "Most Downloaded";
-
-  if (search) {
-    searchParams["query"] = search;
-  } else {
-    // sort: "Highest Rated",
-  }
-
-  const url = new URL(baseUrl);
-  Object.keys(searchParams).forEach((key) =>
-    url.searchParams.append(key, searchParams[key])
-  );
-
-  return url;
-}
-
-export function CivitaiModelRegistry({
-  field,
-}: Pick<AutoFormInputComponentProps, "field">) {
-  const [modelList, setModelList] =
-    React.useState<z.infer<typeof ModelListWrapper>>();
-
-  const [loading, setLoading] = React.useState(false);
-
-  const handleSearch = useDebouncedCallback((search) => {
-    console.log(`Searching... ${search}`);
-
-    setLoading(true);
-
-    const controller = new AbortController();
-    fetch(getUrl(search), {
-      signal: controller.signal,
-    })
-      .then((x) => x.json())
-      .then((a) => {
-        const list = CivitalModelSchema.parse(a);
-        console.log(a);
-
-        setModelList(mapModelsList(list));
-        setLoading(false);
-      });
-
-    return () => {
-      controller.abort();
-      setLoading(false);
-    };
-  }, 300);
-
-  React.useEffect(() => {
-    const controller = new AbortController();
-    fetch(getUrl(), {
-      signal: controller.signal,
-    })
-      .then((x) => x.json())
-      .then((a) => {
-        const list = CivitalModelSchema.parse(a);
-        setModelList(mapModelsList(list));
-      });
-
-    return () => {
-      controller.abort();
-    };
-  }, []);
-
-  return (
-    <ModelSelector
-      field={field}
-      modelList={modelList}
-      label="Civitai"
-      onSearch={handleSearch}
-      shouldFilter={false}
-      isLoading={loading}
-    />
-  );
-}
-
-export function ComfyUIManagerModelRegistry({
-  field,
-}: Pick<AutoFormInputComponentProps, "field">) {
-  const [modelList, setModelList] =
-    React.useState<z.infer<typeof ModelListWrapper>>();
-
-  React.useEffect(() => {
-    const controller = new AbortController();
-    fetch(
-      "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/model-list.json",
-      {
-        signal: controller.signal,
-      }
-    )
-      .then((x) => x.json())
-      .then((a) => {
-        setModelList(ModelListWrapper.parse(a));
-      });
-
-    return () => {
-      controller.abort();
-    };
-  }, []);
-
-  return (
-    <ModelSelector
-      field={field}
-      modelList={modelList}
-      label="ComfyUI Manager"
-    />
-  );
-}
-
-export function ModelSelector({
-  field,
-  modelList,
-  label,
-  onSearch,
-  shouldFilter = true,
-  isLoading,
-}: Pick<AutoFormInputComponentProps, "field"> & {
-  modelList?: z.infer<typeof ModelListWrapper>;
-  label: string;
-  onSearch?: (search: string) => void;
-  shouldFilter?: boolean;
-  isLoading?: boolean;
-}) {
-  const value = (field.value as z.infer<typeof ModelList>) ?? [];
-  const [open, setOpen] = React.useState(false);
-
-  function toggleModel(model: z.infer<typeof Model>) {
-    const prevSelectedModels = value;
-    if (
-      prevSelectedModels.some(
-        (selectedModel) =>
-          selectedModel.url + selectedModel.name === model.url + model.name
-      )
-    ) {
-      field.onChange(
-        prevSelectedModels.filter(
-          (selectedModel) =>
-            selectedModel.url + selectedModel.name !== model.url + model.name
-        )
-      );
-    } else {
-      field.onChange([...prevSelectedModels, model]);
-    }
-  }
-
-  const containerRef = useRef<HTMLDivElement>(null);
-
-  return (
-    <div className="" ref={containerRef}>
-      <Popover open={open} onOpenChange={setOpen}>
-        <PopoverTrigger asChild>
-          <Button
-            variant="outline"
-            role="combobox"
-            aria-expanded={open}
-            className="w-full justify-between flex"
-          >
-            Add from {label}
-            <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
-          </Button>
-        </PopoverTrigger>
-        <PopoverContent className="w-[375px] p-0" side="bottom">
-          <Command shouldFilter={shouldFilter}>
-            <CommandInput
-              placeholder="Search models..."
-              className="h-9"
-              onValueChange={onSearch}
-            >
-              {isLoading && <LoadingIcon />}
-            </CommandInput>
-            <CommandEmpty>No models found.</CommandEmpty>
-            <CommandList className="pointer-events-auto">
-              <CommandGroup>
-                {modelList?.models.map((model) => (
-                  <CommandItem
-                    key={model.url + model.name}
-                    value={model.url}
-                    onSelect={() => {
-                      toggleModel(model);
-                    }}
-                  >
-                    {model.name}
-                    <Check
-                      className={cn(
-                        "ml-auto h-4 w-4",
-                        value.some(
-                          (selectedModel) => selectedModel.url === model.url
-                        )
-                          ? "opacity-100"
-                          : "opacity-0"
-                      )}
-                    />
-                  </CommandItem>
-                ))}
-              </CommandGroup>
-            </CommandList>
-          </Command>
-        </PopoverContent>
-      </Popover>
-    </div>
-  );
-}
@@ -0,0 +1,123 @@
+"use client";
+import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
+import { LoadingIcon } from "@/components/LoadingIcon";
+import { Button } from "@/components/ui/button";
+import {
+  Command,
+  CommandEmpty,
+  CommandGroup,
+  CommandInput,
+  CommandItem,
+  CommandList,
+} from "@/components/ui/command";
+import {
+  Popover,
+  PopoverContent,
+  PopoverTrigger,
+} from "@/components/ui/popover";
+import { cn } from "@/lib/utils";
+import { Check, ChevronsUpDown } from "lucide-react";
+import * as React from "react";
+import { useRef } from "react";
+import { z } from "zod";
+import { ModelListWrapper, Model, ModelList } from "./CivitalModelSchema";
+
+export function ModelSelector({
+  field,
+  modelList,
+  label,
+  onSearch,
+  shouldFilter = true,
+  isLoading,
+  selectMultiple = true,
+}: Pick<AutoFormInputComponentProps, "field"> & {
+  modelList?: z.infer<typeof ModelListWrapper>;
+  label: string;
+  onSearch?: (search: string) => void;
+  shouldFilter?: boolean;
+  isLoading?: boolean;
+  selectMultiple?: boolean;
+}) {
+  const value = (field.value as z.infer<typeof ModelList>) ?? [];
+  const [open, setOpen] = React.useState(false);
+
+  function toggleModel(model: z.infer<typeof Model>) {
+    const prevSelectedModels = value;
+    if (
+      prevSelectedModels.some(
+        (selectedModel) =>
+          selectedModel.url + selectedModel.name === model.url + model.name,
+      )
+    ) {
+      field.onChange(
+        prevSelectedModels.filter(
+          (selectedModel) =>
+            selectedModel.url + selectedModel.name !== model.url + model.name,
+        ),
+      );
+    } else {
+      if (!selectMultiple) {
+        field.onChange([model]);
+      } else {
+        field.onChange([...prevSelectedModels, model]);
+      }
+    }
+  }
+
+  const containerRef = useRef<HTMLDivElement>(null);
+
+  return (
+    <div className="" ref={containerRef}>
+      <Popover open={open} onOpenChange={setOpen}>
+        <PopoverTrigger asChild>
+          <Button
+            variant="outline"
+            role="combobox"
+            aria-expanded={open}
+            className="w-full justify-between flex"
+          >
+            Add from {label}
+            <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
+          </Button>
+        </PopoverTrigger>
+        <PopoverContent className="w-[375px] p-0" side="bottom">
+          <Command shouldFilter={shouldFilter}>
+            <CommandInput
+              placeholder="Search models..."
+              className="h-9"
+              onValueChange={onSearch}
+            >
+              {isLoading && <LoadingIcon />}
+            </CommandInput>
+            <CommandEmpty>No models found.</CommandEmpty>
+            <CommandList className="pointer-events-auto">
+              <CommandGroup>
+                {modelList?.models.map((model) => (
+                  <CommandItem
+                    key={model.url + model.name}
+                    value={model.url}
+                    onSelect={() => {
+                      toggleModel(model);
+                    }}
+                  >
+                    {model.name}
+                    <Check
+                      className={cn(
+                        "ml-auto h-4 w-4",
+                        value.some(
+                          (selectedModel) => selectedModel.url === model.url,
+                        )
+                          ? "opacity-100"
+                          : "opacity-0",
+                      )}
+                    />
+                  </CommandItem>
+                ))}
+              </CommandGroup>
+            </CommandList>
+          </Command>
+        </PopoverContent>
+      </Popover>
+    </div>
+  );
+}
@@ -8,7 +8,9 @@ import {
  AccordionItem,
  AccordionTrigger,
 } from "@/components/ui/accordion";
+import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
+import { Card } from "@/components/ui/card";
 import {
  Command,
  CommandEmpty,
@@ -17,6 +19,22 @@ import {
  CommandItem,
  CommandList,
 } from "@/components/ui/command";
+import {
+  Dialog,
+  DialogClose,
+  DialogContent,
+  DialogDescription,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+  DialogTrigger,
+} from "@/components/ui/dialog";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
 import {
  Popover,
  PopoverContent,
@@ -25,11 +43,18 @@ import {
 import { Textarea } from "@/components/ui/textarea";
 import { cn } from "@/lib/utils";
 import { findAllDeployments } from "@/server/curdDeploments";
-import { Check, ChevronsUpDown } from "lucide-react";
+import {
+  Check,
+  ChevronsUpDown,
+  Edit,
+  ExternalLink,
+  FolderInput,
+  MoreVertical,
+  Plus,
+} from "lucide-react";
 import * as React from "react";
-import { toast } from "sonner";
 import useSWR from "swr";
-import { z } from "zod";
+import { getBranchInfo } from "./getBranchInfo";

 export function SnapshotPickerView({
  field,
@@ -39,19 +64,118 @@ export function SnapshotPickerView({
      <AccordionItem value="item-1">
        <AccordionTrigger className="text-sm">Custom Nodes</AccordionTrigger>
        <AccordionContent className="flex gap-2 flex-col px-1">
-          <SnapshotPresetPicker field={field} />
-          <CustomNodesSelector field={field} />
+          <div className="flex flex-wrap gap-2 justify-end">
+            <CustomNodesSelector field={field} />
+            <SnapshotPresetPicker field={field} />
+            <Dialog>
+              <DialogTrigger asChild>
+                <Button variant="secondary" className="w-fit">
+                  Edit <Edit size={14}></Edit>
+                </Button>
+              </DialogTrigger>
+              <DialogContent className="sm:max-w-[600px] h-full max-h-[600px] grid grid-rows-[auto,1fr,auto]">
+                <DialogHeader>
+                  <DialogTitle>Edit custom nodes</DialogTitle>
+                  <DialogDescription>
+                    Make advacne changes to the snapshots
+                  </DialogDescription>
+                </DialogHeader>
+                <Textarea
+                  className="h-full p-2 max-h-[600px] rounded-md text-xs w-full"
+                  value={JSON.stringify(field.value, null, 2)}
+                  onChange={(e) => {
+                    // Update field.onChange to pass the array of selected models
+                    field.onChange(JSON.parse(e.target.value));
+                  }}
+                />
+                <DialogFooter>
+                  <DialogClose>
+                    <Button type="button" variant="secondary">
+                      Close
+                    </Button>
+                  </DialogClose>
+                </DialogFooter>
+              </DialogContent>
+            </Dialog>
+          </div>
          {field.value && (
-            // <ScrollArea className="w-full bg-gray-100 mx-auto max-w-[500px] rounded-lg">
-            <Textarea
-              className="min-h-[150px] max-h-[300px] p-2 rounded-md text-xs w-full"
-              value={JSON.stringify(field.value, null, 2)}
-              onChange={(e) => {
-                // Update field.onChange to pass the array of selected models
-                field.onChange(JSON.parse(e.target.value));
-              }}
-            />
-            // </ScrollArea>
+            <div className="flex gap-2 flex-col">
+              {Object.entries(field.value.git_custom_nodes).map(
+                ([key, item]: [string, any], index) => (
+                  <Card className="p-4 flex gap-4 items-center justify-between">
+                    <div className="flex gap-4 items-center">
+                      <div className="bg-gray-200 aspect-square w-6 h-6 rounded-full text-center">
+                        {index + 1}
+                      </div>
+                      <div>
+                        <a
+                          target="_blank"
+                          href={key}
+                          className="hover:underline flex items-center gap-2"
+                          rel="noreferrer"
+                        >
+                          <ExternalLink size={12} /> {key}
+                        </a>
+                        <div className="text-2xs text-primary/50">
+                          {item.hash}
+                        </div>
+                      </div>
+                    </div>
+
+                    <DropdownMenu>
+                      <DropdownMenuTrigger asChild type="button">
+                        <Button type="button" variant={"ghost"}>
+                          <MoreVertical size={12} />
+                        </Button>
+                      </DropdownMenuTrigger>
+                      <DropdownMenuContent>
+                        <DropdownMenuItem
+                          disabled={key.endsWith("comfyui-deploy")}
+                          // className="opacity-50"
+                          onClick={() => {
+                            const newNodeList = {
+                              ...field.value.git_custom_nodes,
+                            };
+                            delete newNodeList[key];
+                            const nodeList = newNodeList;
+                            const newValue = {
+                              ...field.value,
+                              git_custom_nodes: nodeList,
+                            };
+                            field.onChange(newValue);
+                          }}
+                        >
+                          Delete
+                        </DropdownMenuItem>
+                        <DropdownMenuItem
+                          // className="opacity-50"
+                          onClick={async () => {
+                            const newNodeList = {
+                              ...field.value.git_custom_nodes,
+                            };
+
+                            const branchInfo = await getBranchInfo(key);
+
+                            if (!branchInfo) return;
+
+                            newNodeList[key].hash = branchInfo?.commit.sha;
+
+                            const nodeList = newNodeList;
+                            const newValue = {
+                              ...field.value,
+                              git_custom_nodes: nodeList,
+                            };
+                            field.onChange(newValue);
+                          }}
+                        >
+                          Update
+                        </DropdownMenuItem>
+                      </DropdownMenuContent>
+                    </DropdownMenu>
+                  </Card>
+                ),
+              )}
+            </div>
          )}
        </AccordionContent>
      </AccordionItem>
@@ -65,13 +189,14 @@ function SnapshotPresetPicker({
  const [open, setOpen] = React.useState(false);
  const [selected, setSelected] = React.useState<string | null>(null);

-  const [frameworks, setFramework] = React.useState<
-    {
-      id: string;
-      label: string;
-      value: string;
-    }[]
-  >();
+  const [frameworks, setFramework] =
+    React.useState<
+      {
+        id: string;
+        label: string;
+        value: string;
+      }[]
+    >();

  React.useEffect(() => {
    findAllDeployments().then((a) => {
@@ -108,12 +233,13 @@ function SnapshotPresetPicker({
          variant="outline"
          role="combobox"
          aria-expanded={open}
-          className="w-full justify-between flex"
+          className="w-fit justify-between flex"
        >
-          {selected
+          <FolderInput size={14} />
+          Import
+          {/* {selected
            ? findItem(selected)?.label
-            : "Select snapshot (From deployments)"}
-          <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
+            : "Select snapshot (From deployments)"} */}
        </Button>
      </PopoverTrigger>
      <PopoverContent className="w-[375px] p-0">
@@ -140,7 +266,7 @@ function SnapshotPresetPicker({
                    "ml-auto h-4 w-4",
                    field.value === framework.value
                      ? "opacity-100"
-                      : "opacity-0"
+                      : "opacity-0",
                  )}
                />
              </CommandItem>
@@ -164,24 +290,6 @@ type CustomNodeList = {
  }[];
 };

-const RepoSchema = z.object({
-  default_branch: z.string(),
-});
-
-const BranchInfoSchema = z.object({
-  commit: z.object({
-    sha: z.string(),
-  }),
-});
-
-function extractRepoName(repoUrl: string) {
-  const url = new URL(repoUrl);
-  const pathParts = url.pathname.split("/");
-  const repoName = pathParts[2].replace(".git", "");
-  const author = pathParts[1];
-  return `${author}/${repoName}`;
-}
-
 function CustomNodesSelector({
  field,
 }: Pick<AutoFormInputComponentProps, "field">) {
@@ -199,12 +307,12 @@ function CustomNodesSelector({

  const { data, error, isLoading } = useSWR<CustomNodeList>(
    "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/custom-node-list.json",
-    fetcher
+    fetcher,
  );

  const keys = React.useMemo(
    () => Object.keys(customNodeList),
-    [customNodeList, data]
+    [customNodeList, data],
  );

  function findItem(value: string) {
@@ -213,6 +321,11 @@ function CustomNodesSelector({
    return included;
  }

+  const onChangeRef = React.useRef(field.onChange);
+  React.useEffect(() => {
+    onChangeRef.current = field.onChange;
+  }, [field.onChange]);
+
  return (
    <Popover open={open} onOpenChange={setOpen}>
      <PopoverTrigger asChild>
@@ -220,10 +333,10 @@ function CustomNodesSelector({
          variant="outline"
          role="combobox"
          aria-expanded={open}
-          className="w-full justify-between flex"
+          className="w-fit justify-between flex"
        >
-          Add custom nodes - {keys.length} selected
-          <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
+          <Plus size={14}></Plus> <Badge>{keys.length} </Badge>
+          {/* <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" /> */}
        </Button>
      </PopoverTrigger>
      <PopoverContent className="w-[375px] p-0" side="bottom">
@@ -243,7 +356,7 @@ function CustomNodesSelector({
                        {
                          hash: string;
                          disabled: boolean;
-                          pip?: string[]
+                          pip?: string[];
                        }
                      >;
                      const x = customNodeList;
@@ -253,57 +366,31 @@ function CustomNodesSelector({
                        delete newNodeList[currentValue];
                        nodeList = newNodeList;
                      } else {
-                        const repoName = extractRepoName(currentValue);
-                        const id = toast.loading(`Fetching repo info...`);
-                        const repo = await fetch(
-                          `https://api.github.com/repos/${repoName}`
-                        )
-                          .then((x) => x.json())
-                          .then((x) => {
-                            console.log(x);
-                            return x;
-                          })
-                          .then((x) => RepoSchema.parse(x))
-                          .catch((e) => {
-                            console.error(e);
-                            toast.dismiss(id);
-                            toast.error(`Failed to fetch repo info ${e.message}`);
-                            return null;
-                          });
-
-                        if (!repo) return;
-                        const branch = repo.default_branch;
-                        const branchInfo = await fetch(
-                          `https://api.github.com/repos/${repoName}/branches/${branch}`
-                        )
-                          .then((x) => x.json())
-                          .then((x) => BranchInfoSchema.parse(x))
-                          .catch((e) => {
-                            console.error(e);
-                            toast.dismiss(id);
-                            toast.error(
-                              `Failed to fetch branch info ${e.message}`
-                            );
-                            return null;
-                          });
-
-                        toast.dismiss(id);
+                        const branchInfo = await getBranchInfo(currentValue);

                        if (!branchInfo) return;

+                        const value: Record<string, any> = {
+                          hash: branchInfo?.commit.sha,
+                          disabled: false,
+                        };
+
+                        if (framework.pip) {
+                          value["pip"] = framework.pip;
+                        }
+
                        nodeList = {
-                          [currentValue]: {
-                            hash: branchInfo?.commit.sha,
-                            disabled: false,
-                            pip: framework.pip
-                          },
                          ...x,
+                          [currentValue]: value,
                        };
                      }
-                      field.onChange({
+
+                      const newValue = {
                        ...field.value,
                        git_custom_nodes: nodeList,
-                      });
+                      };
+
+                      field.onChange(newValue);
                    }}
                  >
                    {framework.title}
@@ -312,7 +399,7 @@ function CustomNodesSelector({
                        "ml-auto h-4 w-4",
                        findItem(framework.reference)
                          ? "opacity-100"
-                          : "opacity-0"
+                          : "opacity-0",
                      )}
                    />
                  </CommandItem>
@@ -0,0 +1,52 @@
+import { toast } from "sonner";
+import { z } from "zod";
+
+const RepoSchema = z.object({
+  default_branch: z.string(),
+});
+const BranchInfoSchema = z.object({
+  commit: z.object({
+    sha: z.string(),
+  }),
+});
+function extractRepoName(repoUrl: string) {
+  const url = new URL(repoUrl);
+  const pathParts = url.pathname.split("/");
+  const repoName = pathParts[2].replace(".git", "");
+  const author = pathParts[1];
+  return `${author}/${repoName}`;
+}
+export async function getBranchInfo(gitUrl: string) {
+  const repoName = extractRepoName(gitUrl);
+  const id = toast.loading(`Fetching repo info...`);
+  const repo = await fetch(`https://api.github.com/repos/${repoName}`)
+    .then((x) => x.json())
+    .then((x) => {
+      console.log(x);
+      return x;
+    })
+    .then((x) => RepoSchema.parse(x))
+    .catch((e) => {
+      console.error(e);
+      toast.dismiss(id);
+      toast.error(`Failed to fetch repo info ${e.message}`);
+      return null;
+    });
+
+  if (!repo) return;
+  const branch = repo.default_branch;
+  const branchInfo = await fetch(
+    `https://api.github.com/repos/${repoName}/branches/${branch}`,
+  )
+    .then((x) => x.json())
+    .then((x) => BranchInfoSchema.parse(x))
+    .catch((e) => {
+      console.error(e);
+      toast.dismiss(id);
+      toast.error(`Failed to fetch branch info ${e.message}`);
+      return null;
+    });
+
+  toast.dismiss(id);
+  return branchInfo;
+}
@@ -0,0 +1,56 @@
+"use client";
+import { z } from "zod";
+import {
+  CivitalModelSchema,
+  ModelListWrapper,
+  Model,
+} from "./CivitalModelSchema";
+
+function mapType(type: string) {
+  switch (type) {
+    case "checkpoint":
+      return "checkpoints";
+  }
+  return type;
+}
+export function mapModelsList(
+  models: z.infer<typeof CivitalModelSchema>,
+): z.infer<typeof ModelListWrapper> {
+  return {
+    models: models.items.flatMap((item) => {
+      return item.modelVersions.map((v) => {
+        return {
+          name: `${item.name} ${v.name} (${v.files[0].name})`,
+          type: mapType(item.type.toLowerCase()),
+          base: v.baseModel,
+          save_path: "default",
+          description: item.description,
+          reference: "",
+          filename: v.files[0].name,
+          // Quick hack to get the download url back as normal url
+          url: `https://civitai.com/models/${v.modelId}?modelVersionId=${v.id}`, //v.files[0].downloadUrl,
+        } as z.infer<typeof Model>;
+      });
+    }),
+  };
+}
+export function getUrl(search?: string) {
+  const baseUrl = "https://civitai.com/api/v1/models";
+  const searchParams = {
+    limit: 5,
+  } as any;
+  searchParams["sort"] = "Most Downloaded";
+
+  if (search) {
+    searchParams["query"] = search;
+  } else {
+    // sort: "Highest Rated",
+  }
+
+  const url = new URL(baseUrl);
+  Object.keys(searchParams).forEach((key) =>
+    url.searchParams.append(key, searchParams[key]),
+  );
+
+  return url;
+}
@@ -0,0 +1,98 @@
+import { AutoFormInputComponentProps } from "@/components/ui/auto-form/types";
+import { getBaseSchema } from "@/components/ui/auto-form/utils";
+import { Badge } from "@/components/ui/badge";
+import {
+  FormItem,
+  FormLabel,
+  FormControl,
+  FormDescription,
+  FormMessage,
+} from "@/components/ui/form";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
+import { Lock } from "lucide-react";
+import * as z from "zod";
+
+export default function AutoFormGPUPicker({
+  label,
+  isRequired,
+  field,
+  fieldConfigItem,
+  zodItem,
+}: AutoFormInputComponentProps) {
+  const baseValues = (getBaseSchema(zodItem) as unknown as z.ZodEnum<any>)._def
+    .values;
+
+  let values: [string, string][] = [];
+  if (!Array.isArray(baseValues)) {
+    values = Object.entries(baseValues);
+  } else {
+    values = baseValues.map((value) => [value, value]);
+  }
+
+  function findItem(value: any) {
+    return values.find((item) => item[0] === value);
+  }
+
+  const plan = fieldConfigItem.inputProps?.sub?.plan;
+  const enabledGPU = ["T4"];
+
+  const planMapping: Record<string, string> = {
+    A10G: "pro",
+    A100: "enterprise",
+  };
+
+  if (plan == "pro") {
+    enabledGPU.push("A10G");
+  } else if (plan == "enterprise") {
+    enabledGPU.push("A10G");
+    enabledGPU.push("A100");
+  }
+
+  return (
+    <FormItem>
+      <FormLabel>
+        {label}
+        {isRequired && <span className="text-destructive"> *</span>}
+      </FormLabel>
+      <FormControl>
+        <Select onValueChange={field.onChange} defaultValue={field.value}>
+          <SelectTrigger>
+            <SelectValue
+              className="w-full"
+              placeholder={fieldConfigItem.inputProps?.placeholder}
+            >
+              {field.value ? findItem(field.value)?.[1] : "Select an option"}
+            </SelectValue>
+          </SelectTrigger>
+          <SelectContent>
+            {values.map(([value, label]) => {
+              const enabled = enabledGPU.includes(value);
+              return (
+                <SelectItem value={label} key={value} disabled={!enabled}>
+                  {label}
+                  {!enabled && (
+                    <span className="mx-2 inline-flex items-center justify-center gap-2">
+                      <Badge className="capitalize">{planMapping[value]}</Badge>{" "}
+                      plan required
+                      <Lock size={14}></Lock>
+                    </span>
+                  )}
+                </SelectItem>
+              );
+            })}
+          </SelectContent>
+        </Select>
+      </FormControl>
+      {fieldConfigItem.description && (
+        <FormDescription>{fieldConfigItem.description}</FormDescription>
+      )}
+      <FormMessage />
+    </FormItem>
+  );
+}
@@ -1,86 +0,0 @@
-// NOTE: this is WIP for doing client side validation for civitai model downloading
-import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
-import { FormControl, FormItem, FormLabel } from "../ui/form";
-import { LoadingIcon } from "@/components/LoadingIcon";
-import * as React from "react";
-import AutoFormInput from "../ui/auto-form/fields/input";
-import { useDebouncedCallback } from "use-debounce";
-import { CivitaiModelResponse } from "@/types/civitai";
-import { z } from "zod";
-import { insertCivitaiModelSchema } from "@/db/schema";
-
-function getUrl(civitai_url: string) {
-  // expect to be a URL to be https://civitai.com/models/36520
-  // possiblity with slugged name and query-param modelVersionId
-
-  const baseUrl = "https://civitai.com/api/v1/models/";
-  const url = new URL(civitai_url);
-  const pathSegments = url.pathname.split("/");
-  const modelId = pathSegments[pathSegments.indexOf("models") + 1];
-  const modelVersionId = url.searchParams.get("modelVersionId");
-
-  return { url: baseUrl + modelId, modelVersionId };
-}
-
-export default function AutoFormCheckpointInput(
-  props: AutoFormInputComponentProps
-) {
-  const [loading, setLoading] = React.useState(false);
-  const [modelRes, setModelRes] =
-    React.useState<z.infer<typeof CivitaiModelResponse>>();
-  const [modelVersionid, setModelVersionId] = React.useState<string | null>();
-  const { label, isRequired, fieldProps, zodItem, fieldConfigItem } = props;
-
-  const handleSearch = useDebouncedCallback((search) => {
-    const validationResult =
-      insertCivitaiModelSchema.shape.civitai_url.safeParse(search);
-    if (!validationResult.success) {
-      console.error(validationResult.error);
-      // Optionally set an error state here
-      return;
-    }
-
-    setLoading(true);
-
-    const controller = new AbortController();
-    const { url, modelVersionId: versionId } = getUrl(search);
-    setModelVersionId(versionId);
-    fetch(url, {
-      signal: controller.signal,
-    })
-      .then((x) => x.json())
-      .then((a) => {
-        const res = CivitaiModelResponse.parse(a);
-        console.log(a);
-        console.log(res);
-        setModelRes(res);
-        setLoading(false);
-      });
-
-    return () => {
-      controller.abort();
-      setLoading(false);
-    };
-  }, 300);
-
-  const modifiedField = {
-    ...fieldProps,
-    // onChange: (event: React.ChangeEvent<HTMLInputElement>) => {
-    //   handleSearch(event.target.value);
-    // },
-  };
-
-  return (
-    <FormItem>
-      {fieldConfigItem.inputProps?.showLabel && (
-        <FormLabel>
-          {label}
-          {isRequired && <span className="text-destructive">*</span>}
-        </FormLabel>
-      )}
-      <FormControl>
-        <AutoFormInput {...props} fieldProps={modifiedField} />
-      </FormControl>
-    </FormItem>
-  );
-}
@@ -0,0 +1,97 @@
+"use client";
+
+import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
+import {
+  FormControl,
+  FormDescription,
+  FormItem,
+  FormLabel,
+  FormMessage,
+} from "../ui/form";
+import { LoadingIcon } from "@/components/LoadingIcon";
+// import { CaretSortIcon, CheckIcon } from "@radix-ui/react-icons";
+import * as React from "react";
+import { Suspense } from "react";
+import {
+  Accordion,
+  AccordionContent,
+  AccordionItem,
+  AccordionTrigger,
+} from "@/components/ui/accordion";
+import { ScrollArea } from "@/components/ui/scroll-area";
+import { Textarea } from "@/components/ui/textarea";
+import { CivitaiModelRegistry } from "./CivitaiModelRegistry";
+import { ComfyUIManagerModelRegistry } from "./ComfyUIManagerModelRegistry";
+import { Input } from "@/components/ui/input";
+import { ModelList } from "@/components/custom-form/CivitalModelSchema";
+import { z } from "zod";
+
+export default function AutoFormModelsPickerUrl({
+  label,
+  isRequired,
+  field,
+  fieldConfigItem,
+  zodItem,
+  fieldProps,
+}: AutoFormInputComponentProps) {
+  return (
+    <FormItem>
+      {fieldConfigItem.inputProps?.showLabel && (
+        <FormLabel>
+          {label}
+          {isRequired && <span className="text-destructive"> *</span>}
+        </FormLabel>
+      )}
+      <FormControl>
+        <Suspense fallback={<LoadingIcon />}>
+          <ModelPickerView field={field} fieldProps={fieldProps} />
+        </Suspense>
+      </FormControl>
+      {fieldConfigItem.description && (
+        <FormDescription>{fieldConfigItem.description}</FormDescription>
+      )}
+      <FormMessage />
+    </FormItem>
+  );
+}
+
+function ModelPickerView({
+  field,
+  fieldProps,
+}: Pick<AutoFormInputComponentProps, "field" | "fieldProps">) {
+  const customOverride = React.useMemo(() => {
+    const customOnChange = (value: z.infer<typeof ModelList>) => {
+      const model = value[0];
+      field.onChange(model?.url);
+    };
+    return {
+      ...field,
+      onChange: customOnChange,
+      value: field.value
+        ? [
+            {
+              url: field.value,
+            },
+          ]
+        : [],
+    };
+  }, [field]);
+
+  return (
+    <div className="flex gap-2 flex-col px-1">
+      <ComfyUIManagerModelRegistry
+        field={customOverride}
+        selectMultiple={false}
+      />
+      <CivitaiModelRegistry field={customOverride} selectMultiple={false} />
+      <Input
+        // className="min-h-[150px] max-h-[300px] p-2 rounded-lg text-xs w-full"
+        value={field.value ?? ""}
+        onChange={(e) => {
+          field.onChange(e.target.value);
+        }}
+        type="text"
+      />
+    </div>
+  );
+}
@@ -1,3 +1,5 @@
+"use client";
+
 import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
 import {
  FormControl,
@@ -7,10 +9,19 @@ import {
  FormMessage,
 } from "../ui/form";
 import { LoadingIcon } from "@/components/LoadingIcon";
-import { ModelPickerView } from "@/components/custom-form/ModelPickerView";
 // import { CaretSortIcon, CheckIcon } from "@radix-ui/react-icons";
 import * as React from "react";
 import { Suspense } from "react";
+import {
+  Accordion,
+  AccordionContent,
+  AccordionItem,
+  AccordionTrigger,
+} from "@/components/ui/accordion";
+import { ScrollArea } from "@/components/ui/scroll-area";
+import { Textarea } from "@/components/ui/textarea";
+import { ComfyUIManagerModelRegistry } from "./ComfyUIManagerModelRegistry";
+import { ExternalLink } from "lucide-react";

 export default function AutoFormModelsPicker({
  label,
@@ -27,6 +38,7 @@ export default function AutoFormModelsPicker({
          {isRequired && <span className="text-destructive"> *</span>}
        </FormLabel>
      )}
+
      <FormControl>
        <Suspense fallback={<LoadingIcon />}>
          <ModelPickerView field={field} />
@@ -35,7 +47,54 @@ export default function AutoFormModelsPicker({
      {fieldConfigItem.description && (
        <FormDescription>{fieldConfigItem.description}</FormDescription>
      )}
+      <FormDescription>
+        {" "}
+        <div className="text-sm">
+          Models are moving to{" "}
+          <a
+            href="/storage"
+            target="_blank"
+            rel="noreferrer"
+            className="inline-flex items-center gap-1 underline"
+          >
+            <ExternalLink size={12} />
+            Storage
+          </a>
+        </div>
+      </FormDescription>
      <FormMessage />
    </FormItem>
  );
 }
+
+function ModelPickerView({
+  field,
+}: Pick<AutoFormInputComponentProps, "field">) {
+  return (
+    <Accordion type="single" collapsible>
+      <AccordionItem value="item-1">
+        <AccordionTrigger className="text-sm">
+          Models (ComfyUI Manager)
+        </AccordionTrigger>
+        <AccordionContent>
+          <div className="flex gap-2 flex-col px-1">
+            <ComfyUIManagerModelRegistry field={field} />
+            {/* <CivitaiModelRegistry field={field} /> */}
+            {/* <span>{field.value.length} selected</span> */}
+            {field.value && (
+              <ScrollArea className="w-full bg-gray-100 mx-auto rounded-lg mt-2">
+                <Textarea
+                  className="min-h-[150px] max-h-[300px] p-2 rounded-lg text-xs w-full"
+                  value={JSON.stringify(field.value, null, 2)}
+                  onChange={(e) => {
+                    field.onChange(JSON.parse(e.target.value));
+                  }}
+                />
+              </ScrollArea>
+            )}
+          </div>
+        </AccordionContent>
+      </AccordionItem>
+    </Accordion>
+  );
+}
@@ -242,6 +242,10 @@ export const navigation: Array<NavGroup> = [
    title: "API",
    links: [{ title: "Endpoints", href: "/docs/endpoints" }],
  },
+  {
+    title: "Video Tutorials",
+    links: [{ title: "Archive", href: "/docs/videos" }],
+  },
 ];

 export function Navigation(props: React.ComponentPropsWithoutRef<"nav">) {
@@ -1,3 +1,4 @@
+import AutoFormGPUPicker from "@/components/custom-form/gpu-picker";
 import AutoFormCheckbox from "./fields/checkbox";
 import AutoFormDate from "./fields/date";
 import AutoFormEnum from "./fields/enum";
@@ -8,6 +9,7 @@ import AutoFormSwitch from "./fields/switch";
 import AutoFormTextarea from "./fields/textarea";
 import AutoFormModelsPicker from "@/components/custom-form/model-picker";
 import AutoFormSnapshotPicker from "@/components/custom-form/snapshot-picker";
+import AutoFormModelsPickerUrl from "@/components/custom-form/model-picker-url-only";

 export const INPUT_COMPONENTS = {
  checkbox: AutoFormCheckbox,
@@ -22,6 +24,8 @@ export const INPUT_COMPONENTS = {
  // Customs
  snapshot: AutoFormSnapshotPicker,
  models: AutoFormModelsPicker,
+  gpuPicker: AutoFormGPUPicker,
+  modelUrlPicker: AutoFormModelsPickerUrl,
 };

 /**
@@ -1,3 +1,4 @@
+import type { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
 import type { INPUT_COMPONENTS } from "./config";
 import type { ControllerRenderProps, FieldValues } from "react-hook-form";
 import type * as z from "zod";
@@ -6,6 +7,7 @@ export type FieldConfigItem = {
  description?: React.ReactNode;
  inputProps?: React.InputHTMLAttributes<HTMLInputElement> & {
    showLabel?: boolean;
+    sub?: Awaited<ReturnType<typeof getCurrentPlanWithAuth>>;
  };
  fieldType?:
    | keyof typeof INPUT_COMPONENTS
@@ -0,0 +1,36 @@
+"use client";
+
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
+import * as React from "react";
+import { createContext, useContext } from "react";
+
+type CurrentPlanContextType = Awaited<
+  ReturnType<typeof getCurrentPlanWithAuth>
+>;
+const CurrentPlanContext = createContext<CurrentPlanContextType | undefined>(
+  undefined,
+);
+
+export function SubscriptionProvider({
+  sub,
+  children,
+}: {
+  sub: CurrentPlanContextType;
+  children: React.ReactNode;
+}) {
+  return (
+    <CurrentPlanContext.Provider value={sub}>
+      {children}
+    </CurrentPlanContext.Provider>
+  );
+}
+
+export const useCurrentPlan = (): CurrentPlanContextType => {
+  const context = useContext(CurrentPlanContext);
+
+  // if (context === undefined) {
+  //   throw new Error("useCurrentPlan must be used within a CurrentPlanProvider");
+  // }
+
+  return context;
+};
@@ -1,3 +1,4 @@
+import { LogsType } from "@/components/LogsViewer";
 import { CivitaiModelResponse } from "@/types/civitai";
 import { type InferSelectModel, relations } from "drizzle-orm";
 import {
@@ -6,10 +7,10 @@ import {
  jsonb,
  pgEnum,
  pgSchema,
+  real,
  text,
  timestamp,
  uuid,
-  real,
 } from "drizzle-orm/pg-core";
 import { createInsertSchema, createSelectSchema } from "drizzle-zod";
 import { TypeOf, z } from "zod";
@@ -104,6 +105,7 @@ export const workflowRunStatus = pgEnum("workflow_run_status", [
  "failed",
  "started",
  "queued",
+  "timeout",
 ]);

 export const deploymentEnvironment = pgEnum("deployment_environment", [
@@ -148,8 +150,9 @@ export const workflowRunsTable = dbSchema.table("workflow_runs", {
      onDelete: "set null",
    },
  ),
-  workflow_inputs:
-    jsonb("workflow_inputs").$type<Record<string, string | number>>(),
+  workflow_inputs: jsonb("workflow_inputs").$type<
+    Record<string, string | number>
+  >(),
  workflow_id: uuid("workflow_id")
    .notNull()
    .references(() => workflowTable.id, {
@@ -172,6 +175,7 @@ export const workflowRunsTable = dbSchema.table("workflow_runs", {
  machine_type: machinesType("machine_type"),
  user_id: text("user_id"),
  org_id: text("org_id"),
+  run_log: jsonb("run_log").$type<LogsType>(),
 });

 export const workflowRunRelations = relations(
@@ -295,8 +299,9 @@ export const deploymentsTable = dbSchema.table("deployments", {
    .references(() => machinesTable.id),
  share_slug: text("share_slug").unique(),
  description: text("description"),
-  showcase_media:
-    jsonb("showcase_media").$type<z.infer<typeof showcaseMedia>>(),
+  showcase_media: jsonb("showcase_media").$type<
+    z.infer<typeof showcaseMedia>
+  >(),
  environment: deploymentEnvironment("environment").notNull(),
  created_at: timestamp("created_at").defaultNow().notNull(),
  updated_at: timestamp("updated_at").defaultNow().notNull(),
@@ -380,12 +385,24 @@ export const resourceUpload = pgEnum("resource_upload", [

 export const modelUploadType = pgEnum("model_upload_type", [
  "civitai",
-  "huggingface",
+  "download-url",
+  "huggingface", // remove?
  "other",
 ]);

 // https://www.answeroverflow.com/m/1125106227387584552
-const modelTypes = ["checkpoint", "lora", "embedding", "vae"] as const;
+export const modelTypes = [
+  "checkpoint",
+  "lora",
+  "embedding",
+  "vae",
+  "clip",
+  "clip_vision",
+  "configs",
+  "controlnet",
+  "upscale_models",
+  "ipadapter",
+] as const;
 export const modelType = pgEnum("model_type", modelTypes);
 export type modelEnumType = (typeof modelTypes)[number];

@@ -399,8 +416,7 @@ export const modelTable = dbSchema.table("models", {
    .notNull()
    .references(() => userVolume.id, {
      onDelete: "cascade",
-    })
-    .notNull(),
+    }),

  model_name: text("model_name"),
  folder_path: text("folder_path"), // in volume
@@ -413,8 +429,10 @@ export const modelTable = dbSchema.table("models", {
    z.infer<typeof CivitaiModelResponse>
  >(),

+  // for our own storage
  hf_url: text("hf_url"),
  s3_url: text("s3_url"),
+
  user_url: text("client_url"),

  is_public: boolean("is_public").notNull().default(true),
@@ -484,16 +502,6 @@ export const subscriptionStatusTable = dbSchema.table("subscription_status", {
  updated_at: timestamp("updated_at").defaultNow().notNull(),
 });

-export const insertCivitaiModelSchema = createInsertSchema(modelTable, {
-  civitai_url: (schema) =>
-    schema.civitai_url
-      .trim()
-      .url({ message: "URL required" })
-      .includes("civitai.com/models", {
-        message: "civitai.com/models link required",
-      }),
-});
-
 export type UserType = InferSelectModel<typeof usersTable>;
 export type WorkflowType = InferSelectModel<typeof workflowTable>;
 export type MachineType = InferSelectModel<typeof machinesTable>;
@@ -31,6 +31,16 @@ const getOutputRoute = createRoute({
                  input_image: "https://somestatic.png",
                },
              }),
+            run_log: (schema) =>
+              schema.run_log.openapi({
+                type: "object",
+                example: [
+                  {
+                    logs: "some logs",
+                    timestamp: 1706631877.3831277,
+                  },
+                ],
+              }),
          }),
        },
      },
@@ -81,7 +91,7 @@ export const registerGetOutputRoute = (app: App) => {
            code: 400,
            message: "Workflow not found",
          },
-          400
+          400,
        );

      return c.json(run, 200);
@@ -94,7 +104,7 @@ export const registerGetOutputRoute = (app: App) => {
        },
        {
          status: 500,
-        }
+        },
      );
    }
  });
@@ -1,5 +1,8 @@
-import { insertCivitaiModelSchema } from "@/db/schema";
+import { z } from "zod";
+import { modelTypes } from "@/db/schema";

-export const addCivitaiModelSchema = insertCivitaiModelSchema.pick({
-  civitai_url: true,
+export const downloadUrlModelSchema = z.object({
+  url: z.string().url(),
+  model_type: z.enum(modelTypes).default("checkpoint")
 });
+
@@ -16,26 +16,14 @@ export const insertCustomMachineSchema = createInsertSchema(machinesTable, {
    schema.snapshot.default({
      comfyui: "d0165d819afe76bd4e6bdd710eb5f3e571b6a804",
      git_custom_nodes: {
-        "https://github.com/BennyKok/comfyui-deploy.git": {
-          hash: "43fe0a384aa5fa9e141d4a264b2ed40a73b817bc",
+        "https://github.com/bennykok/comfyui-deploy": {
+          hash: "df46e3a0e5ad93fa71f5d216997e376af33b2a6d",
          disabled: false,
        },
      },
      file_custom_nodes: [],
    }),
-  models: (schema) =>
-    schema.models.default([
-      {
-        name: "v1-5-pruned-emaonly.ckpt",
-        type: "checkpoints",
-        base: "SD1.5",
-        save_path: "default",
-        description: "Stable Diffusion 1.5 base model",
-        reference: "https://huggingface.co/runwayml/stable-diffusion-v1-5",
-        filename: "v1-5-pruned-emaonly.ckpt",
-        url: "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt",
-      },
-    ]),
+  models: (schema) => schema.models.default([]),
 });

 export const addCustomMachineSchema = insertCustomMachineSchema.pick({
@@ -2,6 +2,7 @@

 import { auth } from "@clerk/nextjs";
 import {
+  modelEnumType,
  modelTable,
  ModelType,
  userVolume,
@@ -10,8 +11,9 @@ import {
 import { withServerPromise } from "./withServerPromise";
 import { db } from "@/db/db";
 import type { z } from "zod";
+import { revalidatePath } from "next/cache";
 import { headers } from "next/headers";
-import { addCivitaiModelSchema } from "./addCivitaiModelSchema";
+import { downloadUrlModelSchema } from "./addCivitaiModelSchema";
 import { and, eq, isNull } from "drizzle-orm";
 import { CivitaiModelResponse, getModelTypeDetails } from "@/types/civitai";

@@ -90,10 +92,10 @@ export async function addModelVolume() {
    .values({
      user_id: userId,
      org_id: orgId,
-      volume_name: `models_${orgId ? orgId: userId}`, // if orgid is avalible use as part of the volume name
-      disabled: false, 
+      volume_name: `models_${orgId ? orgId : userId}`, // if orgid is avalible use as part of the volume name
+      disabled: false,
    })
-    .returning(); 
+    .returning();
  return insertedVolume;
 }

@@ -109,14 +111,191 @@ function getUrl(civitai_url: string) {
  return { url: baseUrl + modelId, modelVersionId };
 }

+// Helper function to make a HEAD request and follow redirects
+async function fetchFinalUrl(
+  url: string,
+): Promise<{ finalUrl: string; dispositionFilename?: string }> {
+  console.log("fetching");
+  const response = await fetch(url, { method: "HEAD", redirect: "follow" });
+  if (!response.ok) {
+    console.log("response not ok");
+    throw new Error(`Request failed with status ${response.status}`);
+  }
+  const contentDisposition = response.headers.get("content-disposition");
+  let filename;
+  if (contentDisposition) {
+    const matches = contentDisposition.match(
+      /filename\*?=['"]?(?:UTF-8'')?([^;'"\n]*)['"]?;?/i,
+    );
+    filename = matches && matches[1]
+      ? decodeURIComponent(matches[1])
+      : undefined;
+  }
+  return { finalUrl: response.url, dispositionFilename: filename };
+}
+
+// The main function for validation
+export const addModel = withServerPromise(
+  async (data: z.infer<typeof downloadUrlModelSchema>) => {
+    const { url } = data;
+
+    if (url.includes("civitai.com/models/")) {
+      // Make a HEAD request to check for 200 OK
+      const response = await fetch(url, { method: "HEAD" });
+      if (!response.ok) {
+        createModelErrorRecord(
+          url,
+          `civitai gave non-ok response`,
+          "civitai",
+          data.model_type,
+        );
+      }
+      addCivitaiModel(data);
+    } else {
+      const { finalUrl, dispositionFilename } = await fetchFinalUrl(url);
+      console.log("finished fetching");
+      console.log(finalUrl, dispositionFilename);
+
+      if (!dispositionFilename) {
+        console.log("no file name");
+        createModelErrorRecord(
+          url,
+          `Could not find a filename from resolved Url: ${finalUrl}`,
+          "download-url",
+          data.model_type,
+        );
+        return;
+      }
+
+      const validExtensions = [".ckpt", ".pt", ".bin", ".pth", ".safetensors"];
+      const extension = dispositionFilename.slice(
+        dispositionFilename.lastIndexOf("."),
+      );
+      if (!validExtensions.includes(extension)) {
+        console.log("invalid extension");
+        createModelErrorRecord(
+          url,
+          `file ext ${extension} is invalid. Valid extensions: ${validExtensions}`,
+          "download-url",
+          data.model_type,
+        );
+      }
+      addModelDownloadUrl(data, dispositionFilename);
+    }
+  },
+);
+
+export const addModelDownloadUrl = withServerPromise(
+  async (data: z.infer<typeof downloadUrlModelSchema>, filename: string) => {
+    console.log("adding model download");
+    const { userId, orgId } = auth();
+    if (!userId) return { error: "No user id" };
+    const volumes = await retrieveModelVolumes();
+
+    const a = await db
+      .insert(modelTable)
+      .values({
+        user_id: userId,
+        org_id: orgId,
+        upload_type: "download-url",
+        model_name: filename,
+        user_url: data.url,
+        user_volume_id: volumes[0].id,
+        model_type: data.model_type,
+      })
+      .returning();
+
+    const b = a[0];
+    console.log("download url about to upload");
+    await uploadModel(data, b, volumes[0]);
+  },
+);
+
+export const deleteModel = withServerPromise(
+  async (modelId: string) => {
+    const model = await db.query.modelTable.findFirst({
+      where: eq(modelTable.id, modelId),
+    });
+
+    // If the model does not exist, throw an error or return a message
+    if (!model) {
+      throw new Error("Model not found");
+      // Or return { error: "Model not found" }; if you prefer to handle it without throwing
+    }
+
+    const volumes = await retrieveModelVolumes();
+    if (
+      model.status === "success" && !!model.folder_path && !!model.model_name
+    ) {
+      const result = await fetch(
+        `${process.env.MODAL_BUILDER_URL!}/delete-volume-model`,
+        {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+          },
+          body: JSON.stringify({
+            volume_name: volumes[0].volume_name,
+            path: model.folder_path,
+            file_name: model.model_name,
+          }),
+        },
+      );
+      if (!result.ok) {
+        const error_log = await result.text();
+        throw new Error(`Error: ${result.statusText} ${error_log}`);
+      }
+    }
+    await db.delete(modelTable).where(eq(modelTable.id, modelId));
+    revalidatePath("/storage");
+    return { message: "Model Deleted" };
+  },
+);
+
+export const getCivitaiModelRes = async (civitaiUrl: string) => {
+  const { url, modelVersionId } = getUrl(civitaiUrl);
+  const civitaiModelRes = await fetch(url)
+    .then((x) => x.json())
+    .then((a) => {
+      return CivitaiModelResponse.parse(a);
+    });
+  return { civitaiModelRes, url, modelVersionId };
+};
+
+const createModelErrorRecord = async (
+  url: string,
+  errorMessage: string,
+  upload_type: "civitai" | "download-url",
+  model_type: modelEnumType,
+) => {
+  const { userId, orgId } = auth();
+  if (!userId) return { error: "No user id" };
+  const volumes = await retrieveModelVolumes();
+
+  const a = await db
+    .insert(modelTable)
+    .values({
+      user_id: userId,
+      org_id: orgId,
+      user_volume_id: volumes[0].id,
+      upload_type: "civitai",
+      model_type,
+      civitai_url: upload_type === "civitai" ? url : undefined,
+      user_url: upload_type === "download-url" ? url : undefined,
+      error_log: errorMessage,
+      status: "failed",
+    })
+    .returning();
+  return a;
+};
+
 export const addCivitaiModel = withServerPromise(
-  async (data: z.infer<typeof addCivitaiModelSchema>) => {
+  async (data: z.infer<typeof downloadUrlModelSchema>) => {
    const { userId, orgId } = auth();

-    if (!data.civitai_url) return { error: "no civitai_url" };
    if (!userId) return { error: "No user id" };

-    const { url, modelVersionId } = getUrl(data?.civitai_url);
+    const { url, modelVersionId } = getUrl(data.url);
    const civitaiModelRes = await fetch(url)
      .then((x) => x.json())
      .then((a) => {
@@ -142,18 +321,17 @@ export const addCivitaiModel = withServerPromise(
      selectedModelVersionId = selectedModelVersion?.id.toString();
    }

-    const userVolume = await getModelVolumes();
-    let cVolume;
-    if (userVolume.length === 0) {
-      const volume = await addModelVolume();
-      cVolume = volume[0];
-    } else {
-      cVolume = userVolume[0];
-    }
+    const volumes = await retrieveModelVolumes();

    const model_type = getModelTypeDetails(civitaiModelRes.type);
    if (!model_type) {
-      return 
+      createModelErrorRecord(
+        url,
+        `Civitai model type ${civitaiModelRes.type} is not currently supported`,
+        "civitai",
+        data.model_type,
+      );
+      return;
    }

    const a = await db
@@ -165,18 +343,18 @@ export const addCivitaiModel = withServerPromise(
        model_name: selectedModelVersion.files[0].name,
        civitai_id: civitaiModelRes.id.toString(),
        civitai_version_id: selectedModelVersionId,
-        civitai_url: data.civitai_url,
-        civitai_download_url: selectedModelVersion.files[0].downloadUrl,
+        civitai_url: data.url,
+        civitai_download_url: selectedModelVersion.files[0].downloadUrl, // there is an issue when a model hoster might put multiple different types of files i.e. their training data.
        civitai_model_response: civitaiModelRes,
-        user_volume_id: cVolume.id,
-        model_type, 
-        updated_at: new Date(),
+        user_volume_id: volumes[0].id,
+        model_type,
      })
      .returning();

    const b = a[0];

-    await uploadModel(data, b, cVolume);
+    await uploadModel(data, b, volumes[0]);
+    revalidatePath("/storage");
  },
 );

@@ -216,7 +394,7 @@ export const addCivitaiModel = withServerPromise(
 // );

 async function uploadModel(
-  data: z.infer<typeof addCivitaiModelSchema>,
+  data: z.infer<typeof downloadUrlModelSchema>,
  c: ModelType,
  v: UserVolumeType,
 ) {
@@ -238,7 +416,9 @@ async function uploadModel(
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
-        download_url: c.civitai_download_url,
+        download_url: c.upload_type === "civitai"
+          ? c.civitai_download_url
+          : c.user_url,
        volume_name: v.volume_name,
        volume_id: v.id,
        model_id: c.id,
@@ -253,7 +433,6 @@ async function uploadModel(
    await db
      .update(modelTable)
      .set({
-        ...data,
        status: "failed",
        error_log: error_log,
      })
@@ -1,5 +1,6 @@
 "use server";

+import { RunDisplay } from "@/components/RunDisplay";
 import { db } from "@/db/db";
 import { deploymentsTable, workflowRunsTable } from "@/db/schema";
 import { count, desc, eq, sql } from "drizzle-orm";
@@ -56,6 +57,15 @@ export async function findAllRuns({
  });
 }

+export async function getAllRunstableContent(props: RunsSearchTypes) {
+  const data = await findAllRunsWithCounts(props);
+
+  return {
+    table: data?.allRuns.map((run) => <RunDisplay run={run} key={run.id} />),
+    total: data?.total,
+  };
+}
+
 export async function findAllRunsWithCounts(props: RunsSearchTypes) {
  const a = await db
    .select({
@@ -1,8 +1,9 @@
 import { db } from "@/db/db";
-import { and, desc, eq, isNull, or } from "drizzle-orm";
+import { and, desc, eq, isNull, ne, or } from "drizzle-orm";
 import { subscriptionStatusTable } from "@/db/schema";
 import { APIKeyUserType } from "@/server/APIKeyBodyRequest";
 import { auth } from "@clerk/nextjs";
+import "server-only";

 export async function getCurrentPlanWithAuth() {
  const { userId, orgId } = auth();
@@ -23,7 +24,11 @@ export async function getCurrentPlan({ user_id, org_id }: APIKeyUserType) {
      eq(subscriptionStatusTable.user_id, user_id),
      org_id
        ? eq(subscriptionStatusTable.org_id, org_id)
-        : or(isNull(subscriptionStatusTable.org_id), eq(subscriptionStatusTable.org_id, "")),
+        : or(
+            isNull(subscriptionStatusTable.org_id),
+            eq(subscriptionStatusTable.org_id, ""),
+          ),
+      ne(subscriptionStatusTable.status, "deleted"),
    ),
    orderBy: desc(subscriptionStatusTable.created_at),
  });
@@ -49,24 +49,26 @@ export async function getRunsData(run_id: string, user?: APIKeyUserType) {
      for (let i = 0; i < data.outputs.length; i++) {
        const output = data.outputs[i];

-        if (output.data?.images !== undefined) {
-          for (let j = 0; j < output.data?.images.length; j++) {
-            const element = output.data?.images[j];
-            element.url = replaceCDNUrl(
-              `${process.env.SPACES_ENDPOINT}/${process.env.SPACES_BUCKET}/outputs/runs/${data.id}/${element.filename}`
-            );
-          }
-        } else if (output.data?.files !== undefined) {
-          for (let j = 0; j < output.data?.files.length; j++) {
-            const element = output.data?.files[j];
-            element.url = replaceCDNUrl(
-              `${process.env.SPACES_ENDPOINT}/${process.env.SPACES_BUCKET}/outputs/runs/${data.id}/${element.filename}`
-            );
-          }
-        }
+        if (output.data?.images !== undefined)
+          replaceUrls(output.data?.images, data.id);
+
+        if (output.data?.files !== undefined)
+          replaceUrls(output.data?.files, data.id);
+
+        if (output.data?.gifs !== undefined)
+          replaceUrls(output.data?.gifs, data.id);
      }
    }
  }

  return data;
 }
+
+function replaceUrls(dataType: any[], dataId: string) {
+  for (let j = 0; j < dataType.length; j++) {
+    const element = dataType[j];
+    element.url = replaceCDNUrl(
+      `${process.env.SPACES_ENDPOINT}/${process.env.SPACES_BUCKET}/outputs/runs/${dataId}/${element.filename}`,
+    );
+  }
+}
@@ -1,16 +1,10 @@
 "use server";

-import { RunOutputs } from "@/components/RunOutputs";
 import { db } from "@/db/db";
 import { workflowRunOutputs } from "@/db/schema";
 import { eq } from "drizzle-orm";

-export async function getRunsOutputDisplay(run_id: string) {
-  return <RunOutputs run_id={run_id} />;
-}
-
 export async function getRunsOutput(run_id: string) {
-  // throw new Error("Not implemented");
  return await db
    .select()
    .from(workflowRunOutputs)
@@ -0,0 +1,18 @@
+"use server";
+
+import type { StringLiteralUnion } from "shikiji";
+import { getHighlighter } from "shikiji";
+
+export async function highlight(
+  code: string,
+  lang: StringLiteralUnion<string>,
+) {
+  const highlighter = await getHighlighter({
+    themes: ["one-dark-pro"],
+    langs: [lang],
+  });
+  return highlighter.codeToHtml(code.trim(), {
+    lang: lang,
+    theme: "one-dark-pro",
+  });
+}
Author	SHA1	Message	Date
bennykok	eb5e2f36cd	chore: update default git url	2024-02-02 00:55:38 +08:00
bennykok	0b0015fbd1	fix: async issues with updating python to 3.11.1	2024-02-02 00:31:04 +08:00
bennykok	abb4a4798c	fix: revert app back with async	2024-02-02 00:11:40 +08:00
bennykok	95238469d6	chore: update custom route	2024-02-01 23:38:35 +08:00
bennykok	86f14b5bce	fix: roll back to old app	2024-02-01 22:11:17 +08:00
bennykok	84802f5c4e	feat: update custom nodes to latest hash	2024-02-01 19:42:46 +08:00
bennykok	68b4d93639	feat(plugin): add run log in comfydeploy plugin level	2024-02-01 18:32:03 +08:00
bennykok	5d349153d2	fix: if timestamp is undefined skip rendering	2024-02-01 18:30:57 +08:00
bennykok	ea675fcd4c	fix: client side refresh for component	2024-02-01 17:01:54 +08:00
bennykok	4c715b815a	feat: bring back refresh indicator	2024-02-01 16:08:50 +08:00
bennykok	bef6ce35de	fix: auto refresh table	2024-02-01 15:24:43 +08:00
bennykok	d8951df35f	fix: log and async task issues with modal script	2024-02-01 13:42:01 +08:00
bennykok	38fea1e79f	fix: docs page	2024-02-01 11:57:15 +08:00
bennykok	66ad3ab4c2	fix: remove base model from new machine settings	2024-01-31 23:33:55 +08:00
bennykok	488d2aee8c	fix: machine build height	2024-01-31 21:41:04 +08:00
bennykok	c4628f6e4c	feat: output render support mp4 and webm	2024-01-31 21:23:00 +08:00
bennykok	9ba349d36a	refactor(plugin): upload files logic in custom routes	2024-01-31 21:22:44 +08:00
bennykok	08ab93127e	chore(plugin): add a 2 secs delay before loading the incoming workflow	2024-01-31 21:22:14 +08:00
bennykok	de641f0acf	fix: log viewer	2024-01-31 21:21:22 +08:00
bennykok	d4d7e98487	fix: comfydeploy npm code display	2024-01-31 18:00:31 +08:00
bennykok	967a77d3a3	fix: gif file not returning file url.	2024-01-31 16:46:08 +08:00
bennykok	a4cd5db360	fix: ensure when getting current pricing plan is not deleted	2024-01-31 16:02:01 +08:00
bennykok	576a6744a4	fix: safe catch stripe changed credential	2024-01-31 15:49:21 +08:00
bennykok	3e5ff7702e	feat: run log dialog	2024-01-31 15:30:29 +08:00
bennykok	009589630d	fix: build	2024-01-31 14:29:42 +08:00
bennykok	8eb2ce3e10	Squashed commit of the following: commit 33c0ad7d14a85f22c57f943dab58610c13d2ac07 Author: Nicholas Koben Kao <[email protected]> Date: Tue Jan 30 21:56:00 2024 -0800 revert custom form change commit d2905ad045ad7856156e3647a81d642999352de7 Merge: 654423d `e3a1d24` Author: Nicholas Koben Kao <[email protected]> Date: Tue Jan 30 20:50:06 2024 -0800 merge schema commit 654423d597e019a5ebf1ab6568c9942fcb9181c5 Author: Nicholas Koben Kao <[email protected]> Date: Tue Jan 30 20:49:34 2024 -0800 merge confl.ict commit 641724c11346319674fbb329e8e29b362117c242 Author: Nicholas Koben Kao <[email protected]> Date: Tue Jan 30 20:47:34 2024 -0800 model reload on create commit eb4dfe8e3f39a0a98eab0fcf1affe7096c12f33b Author: Nicholas Koben Kao <[email protected]> Date: Tue Jan 30 17:00:03 2024 -0800 delete models commit 0bea9583fada102396c4e08fe6da971c94d404df Author: Nicholas Koben Kao <[email protected]> Date: Tue Jan 30 14:35:15 2024 -0800 deploy volume uploader to have timeouts only be modal related	2024-01-31 14:29:36 +08:00
bennykok	e3a1d24304	feat: add auto form value provider	2024-01-31 11:39:13 +08:00
bennykok	2d033570f4	fix: modal logging with complete async io rewrite	2024-01-31 00:41:41 +08:00
bennykok	7ae25aa162	fix: turn log type from text to jsonb	2024-01-31 00:31:00 +08:00
bennykok	4b37de9ec5	fix: turn string check into any for udpate status endpoint	2024-01-30 23:48:15 +08:00
bennykok	b0d1bcc303	feat: add timeout and run_log	2024-01-30 21:42:30 +08:00
bennykok	2193dd287d	fix: ensure logs are disabled by default	2024-01-30 15:17:25 +08:00
bennykok	81bde40aeb	fix: default value	2024-01-30 11:26:50 +08:00
bennykok	317f699c46	tweak	2024-01-30 11:17:41 +08:00
bennykok	0d1bb2aaf4	fix: machine selections	2024-01-30 11:14:46 +08:00
bennykok	b0b23783ba	fix: type	2024-01-30 10:54:21 +08:00
Nicholas Koben Kao	eeb7310955	Merge branch 'doc-update'	2024-01-29 17:47:32 -08:00
Nicholas Koben Kao	e40cc5373f	new docs section	2024-01-29 17:47:23 -08:00
bennykok	debdaf418c	feat: revamp custom nodes install	2024-01-29 22:27:05 +08:00
Nick Kao	90107ebe1b	change to show npm install instead of copying file (#5 )	2024-01-29 22:04:09 +08:00
bennykok	a90c6c1db4	chore: machine list machine type badge	2024-01-29 22:04:09 +08:00
bennykok	1bf3c1dcd0	fix: machine selection might revert when being refreshed	2024-01-29 22:04:06 +08:00
bennykok	24e95a1954	fix: add machine model dialog	2024-01-29 13:38:00 +08:00
bennykok	194715920e	feat: add models redirect	2024-01-29 02:09:56 +08:00
bennykok	0cb1e92f4f	reduce the idle timeout	2024-01-29 01:59:11 +08:00
bennykok	d96811a0c3	fix: civitai picker	2024-01-28 22:08:31 +08:00
bennykok	757c587901	feat: add drop down selection for storage	2024-01-28 20:42:15 +08:00
bennykok	3b7db4480b	update default comfy deploy version	2024-01-28 19:28:40 +08:00
bennykok	10bbb393a7	fix: invalid input for custom nodes	2024-01-28 17:48:22 +08:00
bennykok	c1fc06fd39	ui tweak	2024-01-28 14:41:14 +08:00
bennykok	9de266fbab	chore: add pricing plan lock for machine types	2024-01-28 14:36:36 +08:00
Nick Kaoandbennykok	42aaf1acb9	Nickkao/volume improvemnts v3 (#4 ) * fix: attempt fixing timeout * be validation work * arbitrary model input, BE validation, error record creation with error logs during potential failure points * remove unused type --------- Co-authored-by: bennykok <[email protected]>	2024-01-28 11:09:00 +08:00
Nicholas Koben Kao	852d889397	clearer deploy instruct	2024-01-27 09:58:40 -08:00
bennykok	cb01c896a0	feat: adding add models lock according to prcing plan	2024-01-28 00:51:44 +08:00
bennykok	fbb7b18273	chore: enable promote code	2024-01-28 00:35:19 +08:00
				`@@ -0,0 +1 @@`
				`ALTER TYPE "model_upload_type" ADD VALUE 'download_url';`