From 2cc0670da9a7731cf162c127dffcb2703c255eb1 Mon Sep 17 00:00:00 2001 From: minkhant1996-dev Date: Thu, 31 Oct 2024 17:49:12 +0700 Subject: [PATCH] update --- folder_paths.py | 11 +- min-comfyui-config.yaml | 47 +----- src_min/rp_handler.py | 350 ---------------------------------------- src_min/start.sh | 20 --- src_min/test_input.json | 63 -------- 5 files changed, 10 insertions(+), 481 deletions(-) delete mode 100644 src_min/rp_handler.py delete mode 100644 src_min/start.sh delete mode 100644 src_min/test_input.json diff --git a/folder_paths.py b/folder_paths.py index 7801cbf70..6afbe369d 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -6,13 +6,20 @@ import mimetypes import logging from typing import Set, List, Dict, Tuple, Literal from collections.abc import Collection - +import yaml supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl', '.sft'} folder_names_and_paths: dict[str, tuple[list[str], set[str]]] = {} base_path = os.path.dirname(os.path.realpath(__file__)) -models_dir = "/home/minkhant/Documents/BrookAI/AI_MODELS" #os.path.join(base_path, "models") +with open(os.path.join(base_path, "min-comfyui-config.yaml"), "r") as f: + config = yaml.safe_load(f) + if "models_folder" in config: + models_folder = config["models_folder"] + else: + models_folder = "models" +models_dir = os.path.join("/".join(base_path.split("/")[:-1]), models_folder) +print("models_dir", models_dir) folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions) folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"]) folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions) diff --git a/min-comfyui-config.yaml b/min-comfyui-config.yaml index ff65dc9ec..54e8dfc7a 100644 --- a/min-comfyui-config.yaml +++ b/min-comfyui-config.yaml @@ -1,47 +1,2 @@ -#Rename this to extra_model_paths.yaml and ComfyUI will load it - -#config for a1111 ui -#all you have to do is change the base_path to where yours is installed -a111: - base_path: path/to/stable-diffusion-webui/ - - checkpoints: models/Stable-diffusion - configs: models/Stable-diffusion - vae: models/VAE - loras: | - models/Lora - models/LyCORIS - upscale_models: | - models/ESRGAN - models/RealESRGAN - models/SwinIR - embeddings: embeddings - hypernetworks: models/hypernetworks - controlnet: models/ControlNet - -#config for comfyui -#your base path should be either an existing comfy install or a central folder where you store all of your models, loras, etc. - -comfyui: - base_path: "/home/minkhant/Documents/BrookAI/AI_MODELS" - # You can use is_default to mark that these folders should be listed first, and used as the default dirs for eg downloads - #is_default: true - checkpoints: models/checkpoints/ - clip: models/clip/ - clip_vision: models/clip_vision/ - configs: models/configs/ - controlnet: models/controlnet/ - diffusion_models: | - models/diffusion_models - models/unet - embeddings: models/embeddings/ - loras: models/loras/ - upscale_models: models/upscale_models/ - vae: models/vae/ - -#other_ui: -# base_path: path/to/ui -# checkpoints: models/checkpoints -# gligen: models/gligen -# custom_nodes: path/custom_nodes +models_folder: /home/minkhant/Documents/BrookAI/AI_MODELS \ No newline at end of file diff --git a/src_min/rp_handler.py b/src_min/rp_handler.py deleted file mode 100644 index 2150284a6..000000000 --- a/src_min/rp_handler.py +++ /dev/null @@ -1,350 +0,0 @@ -import runpod -from runpod.serverless.utils import rp_upload -import json -import urllib.request -import urllib.parse -import time -import os -import requests -import base64 -from io import BytesIO - -# 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 -# Time to wait between poll attempts in milliseconds -COMFY_POLLING_INTERVAL_MS = os.environ.get("COMFY_POLLING_INTERVAL_MS", 250) -# Maximum number of poll attempts -COMFY_POLLING_MAX_RETRIES = os.environ.get("COMFY_POLLING_MAX_RETRIES", 500) -# Host where ComfyUI is running -COMFY_HOST = "127.0.0.1:8188" -# Enforce a clean state after each job is done -# see https://docs.runpod.io/docs/handler-additional-controls#refresh-worker -REFRESH_WORKER = os.environ.get("REFRESH_WORKER", "false").lower() == "true" - - -def validate_input(job_input): - """ - Validates the input for the handler function. - - Args: - job_input (dict): The input data to validate. - - Returns: - tuple: A tuple containing the validated data and an error message, if any. - The structure is (validated_data, error_message). - """ - # Validate if job_input is provided - if job_input is None: - return None, "Please provide input" - - # Check if input is a string and try to parse it as JSON - if isinstance(job_input, str): - try: - job_input = json.loads(job_input) - except json.JSONDecodeError: - return None, "Invalid JSON format in input" - - # Validate 'workflow' in input - workflow = job_input.get("workflow") - if workflow is None: - return None, "Missing 'workflow' parameter" - - # Validate 'images' in input, if provided - images = job_input.get("images") - if images is not None: - if not isinstance(images, list) or not all( - "name" in image and "image" in image for image in images - ): - return ( - None, - "'images' must be a list of objects with 'name' and 'image' keys", - ) - - # Return validated data and no error - return {"workflow": workflow, "images": images}, None - - -def check_server(url, retries=500, delay=50): - """ - 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): - 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: - # If an exception occurs, the server may not be ready - pass - - # Wait for the specified delay before retrying - time.sleep(delay / 1000) - - print( - f"runpod-worker-comfy - Failed to connect to server at {url} after {retries} attempts." - ) - return False - - -def upload_images(images): - """ - Upload a list of base64 encoded images to the ComfyUI server using the /upload/image endpoint. - - Args: - images (list): A list of dictionaries, each containing the 'name' of the image and the 'image' as a base64 encoded string. - server_address (str): The address of the ComfyUI server. - - Returns: - list: A list of responses from the server for each image upload. - """ - if not images: - return {"status": "success", "message": "No images to upload", "details": []} - - responses = [] - upload_errors = [] - - print(f"runpod-worker-comfy - image(s) upload") - - for image in images: - name = image["name"] - image_data = image["image"] - blob = base64.b64decode(image_data) - - # Prepare the form data - files = { - "image": (name, BytesIO(blob), "image/png"), - "overwrite": (None, "true"), - } - - # POST request to upload the image - response = requests.post(f"http://{COMFY_HOST}/upload/image", files=files) - if response.status_code != 200: - upload_errors.append(f"Error uploading {name}: {response.text}") - else: - responses.append(f"Successfully uploaded {name}") - - if upload_errors: - print(f"runpod-worker-comfy - image(s) upload with errors") - return { - "status": "error", - "message": "Some images failed to upload", - "details": upload_errors, - } - - print(f"runpod-worker-comfy - image(s) upload complete") - return { - "status": "success", - "message": "All images uploaded successfully", - "details": responses, - } - - -def queue_workflow(workflow): - """ - Queue a workflow to be processed by ComfyUI - - Args: - workflow (dict): A dictionary containing the workflow to be processed - - Returns: - dict: The JSON response from ComfyUI after processing the workflow - """ - - # The top level element "prompt" is required by ComfyUI - data = json.dumps({"prompt": workflow}).encode("utf-8") - - req = urllib.request.Request(f"http://{COMFY_HOST}/prompt", data=data) - return json.loads(urllib.request.urlopen(req).read()) - - -def get_history(prompt_id): - """ - Retrieve the history of a given prompt using its ID - - Args: - prompt_id (str): The ID of the prompt whose history is to be retrieved - - Returns: - dict: The history of the prompt, containing all the processing steps and results - """ - with urllib.request.urlopen(f"http://{COMFY_HOST}/history/{prompt_id}") as response: - return json.loads(response.read()) - - -def base64_encode(img_path): - """ - Returns base64 encoded image. - - Args: - img_path (str): The path to the image - - Returns: - str: The base64 encoded image - """ - with open(img_path, "rb") as image_file: - encoded_string = base64.b64encode(image_file.read()).decode("utf-8") - return f"{encoded_string}" - - -def process_output_images(outputs, job_id): - """ - This function takes the "outputs" from image generation and the job ID, - then determines the correct way to return the image, either as a direct URL - to an AWS S3 bucket or as a base64 encoded string, depending on the - environment configuration. - - Args: - outputs (dict): A dictionary containing the outputs from image generation, - typically includes node IDs and their respective output data. - job_id (str): The unique identifier for the job. - - Returns: - dict: A dictionary with the status ('success' or 'error') and the message, - which is either the URL to the image in the AWS S3 bucket or a base64 - encoded string of the image. In case of error, the message details the issue. - - The function works as follows: - - It first determines the output path for the images from an environment variable, - defaulting to "/comfyui/output" if not set. - - It then iterates through the outputs to find the filenames of the generated images. - - After confirming the existence of the image in the output folder, it checks if the - AWS S3 bucket is configured via the BUCKET_ENDPOINT_URL environment variable. - - If AWS S3 is configured, it uploads the image to the bucket and returns the URL. - - If AWS S3 is not configured, it encodes the image in base64 and returns the string. - - If the image file does not exist in the output folder, it returns an error status - with a message indicating the missing image file. - """ - - # The path where ComfyUI stores the generated images - COMFY_OUTPUT_PATH = os.environ.get("COMFY_OUTPUT_PATH", "/comfyui/output") - - output_images = {} - - for node_id, node_output in outputs.items(): - if "images" in node_output: - for image in node_output["images"]: - output_images = os.path.join(image["subfolder"], image["filename"]) - - print(f"runpod-worker-comfy - image generation is done") - - # expected image output folder - local_image_path = f"{COMFY_OUTPUT_PATH}/{output_images}" - - print(f"runpod-worker-comfy - {local_image_path}") - - # The image is in the output folder - if os.path.exists(local_image_path): - if os.environ.get("BUCKET_ENDPOINT_URL", False): - # URL to image in AWS S3 - image = rp_upload.upload_image(job_id, local_image_path) - print( - "runpod-worker-comfy - the image was generated and uploaded to AWS S3" - ) - else: - # base64 image - image = base64_encode(local_image_path) - print( - "runpod-worker-comfy - the image was generated and converted to base64" - ) - - return { - "status": "success", - "message": image, - } - else: - print("runpod-worker-comfy - the image does not exist in the output folder") - return { - "status": "error", - "message": f"the image does not exist in the specified output folder: {local_image_path}", - } - - -def handler(job): - """ - The main function that handles a job of generating an image. - - This function validates the input, sends a prompt to ComfyUI for processing, - polls ComfyUI for result, and retrieves generated images. - - Args: - job (dict): A dictionary containing job details and input parameters. - - Returns: - dict: A dictionary containing either an error message or a success status with generated images. - """ - job_input = job["input"] - - # Make sure that the input is valid - validated_data, error_message = validate_input(job_input) - if error_message: - return {"error": error_message} - - # Extract validated data - workflow = validated_data["workflow"] - images = validated_data.get("images") - - # Make sure that the ComfyUI API is available - check_server( - f"http://{COMFY_HOST}", - COMFY_API_AVAILABLE_MAX_RETRIES, - COMFY_API_AVAILABLE_INTERVAL_MS, - ) - - # Upload images if they exist - upload_result = upload_images(images) - - if upload_result["status"] == "error": - return upload_result - - # Queue the workflow - try: - queued_workflow = queue_workflow(workflow) - prompt_id = queued_workflow["prompt_id"] - print(f"runpod-worker-comfy - queued workflow with ID {prompt_id}") - except Exception as e: - return {"error": f"Error queuing workflow: {str(e)}"} - - # Poll for completion - print(f"runpod-worker-comfy - wait until image generation is complete") - retries = 0 - try: - while retries < COMFY_POLLING_MAX_RETRIES: - history = get_history(prompt_id) - - # Exit the loop if we have found the history - if prompt_id in history and history[prompt_id].get("outputs"): - 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)}"} - - # Get the generated image and return it as URL in an AWS bucket or as base64 - images_result = process_output_images(history[prompt_id].get("outputs"), job["id"]) - - result = {**images_result, "refresh_worker": REFRESH_WORKER} - - return result - - -# Start the handler only if this script is run directly -if __name__ == "__main__": - runpod.serverless.start({"handler": handler}) diff --git a/src_min/start.sh b/src_min/start.sh deleted file mode 100644 index 6cb5cb84a..000000000 --- a/src_min/start.sh +++ /dev/null @@ -1,20 +0,0 @@ -#!/usr/bin/env bash - -# Use libtcmalloc for better memory management -TCMALLOC="$(ldconfig -p | grep -Po "libtcmalloc.so.\d" | head -n 1)" -export LD_PRELOAD="${TCMALLOC}" - -# Serve the API and don't shutdown the container -if [ "$SERVE_API_LOCALLY" == "true" ]; then - echo "runpod-worker-comfy: Starting ComfyUI" - python3 /comfyui/main.py --disable-auto-launch --disable-metadata --listen & - - echo "runpod-worker-comfy: Starting RunPod Handler" - python3 -u /rp_handler.py --rp_serve_api --rp_api_host=0.0.0.0 -else - echo "runpod-worker-comfy: Starting ComfyUI" - python3 /comfyui/main.py --disable-auto-launch --disable-metadata & - - echo "runpod-worker-comfy: Starting RunPod Handler" - python3 -u /rp_handler.py -fi \ No newline at end of file diff --git a/src_min/test_input.json b/src_min/test_input.json deleted file mode 100644 index 7fdb0f200..000000000 --- a/src_min/test_input.json +++ /dev/null @@ -1,63 +0,0 @@ -{ - "input": { - "workflow": { - "3": { - "inputs": { - "seed": 234234, - "steps": 20, - "cfg": 8, - "sampler_name": "euler", - "scheduler": "normal", - "denoise": 1, - "model": ["4", 0], - "positive": ["6", 0], - "negative": ["7", 0], - "latent_image": ["5", 0] - }, - "class_type": "KSampler" - }, - "4": { - "inputs": { - "ckpt_name": "sd_xl_base_1.0.safetensors" - }, - "class_type": "CheckpointLoaderSimple" - }, - "5": { - "inputs": { - "width": 512, - "height": 512, - "batch_size": 1 - }, - "class_type": "EmptyLatentImage" - }, - "6": { - "inputs": { - "text": "beautiful scenery nature glass bottle landscape, purple galaxy bottle,", - "clip": ["4", 1] - }, - "class_type": "CLIPTextEncode" - }, - "7": { - "inputs": { - "text": "text, watermark", - "clip": ["4", 1] - }, - "class_type": "CLIPTextEncode" - }, - "8": { - "inputs": { - "samples": ["3", 0], - "vae": ["4", 2] - }, - "class_type": "VAEDecode" - }, - "9": { - "inputs": { - "filename_prefix": "ComfyUI/test", - "images": ["8", 0] - }, - "class_type": "SaveImage" - } - } - } -}