feat: add millon js, add models picker dialog, update builder

This commit is contained in:
BennyKok
2024-01-07 17:22:28 +08:00
parent 3c4bce630e
commit 01a9c1a1d6
33 changed files with 1693 additions and 190 deletions
+125 -58
View File
@@ -1,9 +1,10 @@
from typing import Union, Optional, Dict
from pydantic import BaseModel
from typing import Union, Optional, Dict, List
from pydantic import BaseModel, Field, field_validator
from fastapi import FastAPI, HTTPException, WebSocket, BackgroundTasks, WebSocketDisconnect
from fastapi.responses import JSONResponse
from fastapi.logger import logger as fastapi_logger
import os
from enum import Enum
import json
import subprocess
import time
@@ -104,17 +105,49 @@ class Snapshot(BaseModel):
comfyui: str
git_custom_nodes: Dict[str, GitCustomNodes]
class Model(BaseModel):
name: str
type: str
base: str
save_path: str
description: str
reference: str
filename: str
url: str
class GPUType(str, Enum):
T4 = "T4"
A10G = "A10G"
A100 = "A100"
L4 = "L4"
class Item(BaseModel):
machine_id: str
name: str
snapshot: Snapshot
models: List[Model]
callback_url: str
gpu: GPUType = Field(default=GPUType.T4)
@field_validator('gpu')
@classmethod
def check_gpu(cls, value):
if value not in GPUType.__members__:
raise ValueError(f"Invalid GPU option. Choose from: {', '.join(GPUType.__members__.keys())}")
return GPUType(value)
@app.websocket("/ws/{machine_id}")
async def websocket_endpoint(websocket: WebSocket, machine_id: str):
await websocket.accept()
machine_id_websocket_dict[machine_id] = websocket
# Send existing logs
if machine_id in machine_logs_cache:
await websocket.send_text(json.dumps({"event": "LOGS", "data": {
"machine_id": machine_id,
"logs": json.dumps(machine_logs_cache[machine_id]) ,
"timestamp": time.time()
}}))
try:
while True:
data = await websocket.receive_text()
@@ -156,6 +189,9 @@ async def create_item(item: Item):
return JSONResponse(status_code=200, content={"message": "Build Queued"})
# Initialize the logs cache
machine_logs_cache = {}
async def build_logic(item: Item):
# Deploy to modal
folder_path = f"/app/builds/{item.machine_id}"
@@ -175,7 +211,8 @@ async def build_logic(item: Item):
# Write the config file
config = {
"name": item.name,
"deploy_test": os.environ.get("DEPLOY_TEST_FLAG", "False")
"deploy_test": os.environ.get("DEPLOY_TEST_FLAG", "False"),
"gpu": item.gpu
}
with open(f"{folder_path}/config.py", "w") as f:
f.write("config = " + json.dumps(config))
@@ -183,79 +220,99 @@ async def build_logic(item: Item):
with open(f"{folder_path}/data/snapshot.json", "w") as f:
f.write(item.snapshot.json())
with open(f"{folder_path}/data/models.json", "w") as f:
models_json_list = [model.dict() for model in item.models]
models_json_string = json.dumps(models_json_list)
f.write(models_json_string)
# os.chdir(folder_path)
# process = subprocess.Popen(f"modal deploy {folder_path}/app.py", stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True)
process = await asyncio.subprocess.create_subprocess_shell(
f"modal deploy app.py",
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=folder_path
cwd=folder_path,
# env={**os.environ, "PYTHONUNBUFFERED": "1"}
)
url = None
# Initialize the logs cache
machine_logs_cache = []
# Stream the output
# Read output
while True:
line = await process.stdout.readline()
error = await process.stderr.readline()
if not line and not error:
break
l = line.decode('utf-8').strip()
e = error.decode('utf-8').strip()
if item.machine_id not in machine_logs_cache:
machine_logs_cache[item.machine_id] = []
if l != "":
logger.info(l)
machine_logs_cache.append({
"logs": l,
"timestamp": time.time()
})
machine_logs = machine_logs_cache[item.machine_id]
if item.machine_id in machine_id_websocket_dict:
await machine_id_websocket_dict[item.machine_id].send_text(json.dumps({"event": "LOGS", "data": {
"machine_id": item.machine_id,
"logs": l,
"timestamp": time.time()
}}))
async def read_stream(stream, isStderr):
while True:
line = await stream.readline()
if line:
l = line.decode('utf-8').strip()
if l == "":
continue
if "Created comfyui_app =>" in l or (l.startswith("https://") and l.endswith(".modal.run")):
if "Created comfyui_app =>" in l:
url = l.split("=>")[1].strip()
else:
# Some case it only prints the url on a blank line
url = l
if url:
machine_logs_cache.append({
"logs": f"App image built, url: {url}",
if not isStderr:
logger.info(l)
machine_logs.append({
"logs": l,
"timestamp": time.time()
})
if item.machine_id in machine_id_websocket_dict:
await machine_id_websocket_dict[item.machine_id].send_text(json.dumps({"event": "LOGS", "data": {
"machine_id": item.machine_id,
"logs": f"App image built, url: {url}",
"logs": l,
"timestamp": time.time()
}}))
if e != "":
logger.info(e)
machine_logs_cache.append({
"logs": e,
"timestamp": time.time()
})
if item.machine_id in machine_id_websocket_dict:
await machine_id_websocket_dict[item.machine_id].send_text(json.dumps({"event": "LOGS", "data": {
"machine_id": item.machine_id,
"logs": e,
"timestamp": time.time()
}}))
if "Created comfyui_app =>" in l or (l.startswith("https://") and l.endswith(".modal.run")):
if "Created comfyui_app =>" in l:
url = l.split("=>")[1].strip()
else:
# Some case it only prints the url on a blank line
url = l
if url:
machine_logs.append({
"logs": f"App image built, url: {url}",
"timestamp": time.time()
})
if item.machine_id in machine_id_websocket_dict:
await machine_id_websocket_dict[item.machine_id].send_text(json.dumps({"event": "LOGS", "data": {
"machine_id": item.machine_id,
"logs": f"App image built, url: {url}",
"timestamp": time.time()
}}))
await machine_id_websocket_dict[item.machine_id].send_text(json.dumps({"event": "FINISHED", "data": {
"status": "succuss",
}}))
else:
# is error
logger.error(l)
machine_logs.append({
"logs": e,
"timestamp": time.time()
})
if item.machine_id in machine_id_websocket_dict:
await machine_id_websocket_dict[item.machine_id].send_text(json.dumps({"event": "LOGS", "data": {
"machine_id": item.machine_id,
"logs": e,
"timestamp": time.time()
}}))
await machine_id_websocket_dict[item.machine_id].send_text(json.dumps({"event": "FINISHED", "data": {
"status": "failed",
}}))
else:
break
stdout_task = asyncio.create_task(read_stream(process.stdout, False))
stderr_task = asyncio.create_task(read_stream(process.stderr, True))
await asyncio.wait([stdout_task, stderr_task])
# Wait for the subprocess to finish
await process.wait()
@@ -273,29 +330,39 @@ async def build_logic(item: Item):
if process.returncode != 0:
logger.info("An error occurred.")
# Send a post request with the json body machine_id to the callback url
machine_logs_cache.append({
machine_logs.append({
"logs": "Unable to build the app image.",
"timestamp": time.time()
})
requests.post(item.callback_url, json={"machine_id": item.machine_id, "build_log": json.dumps(machine_logs_cache)})
requests.post(item.callback_url, json={"machine_id": item.machine_id, "build_log": json.dumps(machine_logs)})
if item.machine_id in machine_logs_cache:
del machine_logs_cache[item.machine_id]
return
# return JSONResponse(status_code=400, content={"error": "Unable to build the app image."})
# app_suffix = "comfyui-app"
if url is None:
machine_logs_cache.append({
machine_logs.append({
"logs": "App image built, but url is None, unable to parse the url.",
"timestamp": time.time()
})
requests.post(item.callback_url, json={"machine_id": item.machine_id, "build_log": json.dumps(machine_logs_cache)})
requests.post(item.callback_url, json={"machine_id": item.machine_id, "build_log": json.dumps(machine_logs)})
if item.machine_id in machine_logs_cache:
del machine_logs_cache[item.machine_id]
return
# return JSONResponse(status_code=400, content={"error": "App image built, but url is None, unable to parse the url."})
# example https://bennykok--my-app-comfyui-app.modal.run/
# my_url = f"https://{MODAL_ORG}--{item.container_id}-{app_suffix}.modal.run"
requests.post(item.callback_url, json={"machine_id": item.machine_id, "endpoint": url, "build_log": json.dumps(machine_logs_cache)})
requests.post(item.callback_url, json={"machine_id": item.machine_id, "endpoint": url, "build_log": json.dumps(machine_logs)})
if item.machine_id in machine_logs_cache:
del machine_logs_cache[item.machine_id]
logger.info("done")
logger.info(url)
@@ -57,6 +57,7 @@ WORKDIR /
COPY /data/install_deps.py .
COPY /data/deps.json .
COPY /data/models.json .
RUN python3 install_deps.py
WORKDIR /comfyui/custom_nodes
+1 -1
View File
@@ -105,7 +105,7 @@ image = Image.debian_slim()
target_image = image if deploy_test else dockerfile_image
@stub.function(image=target_image, gpu="T4")
@stub.function(image=target_image, gpu=config["gpu"])
def run(input: Input):
import subprocess
import time
+1 -1
View File
@@ -1 +1 @@
config = {"name": "my-app", "deploy_test": "True"}
config = {"name": "my-app", "deploy_test": "True", "gpu": "T4"}
@@ -28,9 +28,10 @@ def check_server(url, retries=50, delay=500):
)
return False
check_server("http://127.0.0.1:8188")
root_url = "http://127.0.0.1:8188"
check_server(root_url)
url = "http://127.0.0.1:8188/customnode/install"
headers = {"Content-Type": "application/json"}
# Load JSON array from deps.json
@@ -39,12 +40,15 @@ with open('deps.json') as f:
# Make a POST request for each package
for package in packages:
response = requests.request("POST", url, json=package, headers=headers)
response = requests.request("POST", f"{root_url}/customnode/install", json=package, headers=headers)
print(response.text)
# restore_snapshot_url = "http://127.0.0.1:8188/snapshot/restore?target=snapshot"
# response = requests.request("GET", restore_snapshot_url, headers=headers)
# print(response.text)
with open('models.json') as f:
models = json.load(f)
for model in models:
response = requests.request("POST", f"{root_url}/model/install", json=model, headers=headers)
print(response.text)
# Close the server
server_process.terminate()
@@ -0,0 +1,12 @@
[
{
"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"
}
]