Upload run_comfyui.py

run_comfyui.py

@@ -0,0 +1,255 @@
+import requests
+import json
+import time
+import argparse
+import os
+import re
+
+def check_comfyui_started(url):
+    """
+    Checks if ComfyUI is running at the given URL.
+    """
+    try:
+        response = requests.get(url)
+        return response.status_code == 200
+    except requests.ConnectionError:
+        return False
+
+def load_and_update_prompt(file_path, width, height, length, batch_size, shift_amount, guidance_amount, steps_amount, denoise_amount, frame_rate, prompt):
+    """
+    Loads the prompt payload from a JSON file and updates the specified fields.
+
+    Args:
+        file_path (str): Path to the prompt-payload.json file.
+        width (int): Width to set in the prompt.
+        height (int): Height to set in the prompt.
+        length (int): Length (frames) to set in the prompt.
+        batch_size (int): Batch size to set in the prompt.
+        shift_amount (float): Shift amount to set in the prompt.
+        guidance_amount (float): Guidance amount to set in the prompt.
+        steps_amount (int): Steps amount to set in the prompt.
+        denoise_amount (float): Denoise amount to set in the prompt.
+        frame_rate (int): Frame rate to set in the prompt.
+        prompt (str): Prompt text to set in the prompt.
+
+    Returns:
+        dict: Updated prompt payload as a dictionary.
+    """
+    try:
+        with open(file_path, 'r', encoding='utf-8') as f:
+            payload = json.load(f)
+    except Exception as e:
+        print(f"Error loading JSON file: {e}")
+        return None
+
+    # Update width and height
+    payload['prompt']['232']['inputs']['width'] = width
+    payload['prompt']['232']['inputs']['height'] = height
+
+    # Update length (frames)
+    payload['prompt']['295']['inputs']['int'] = length
+
+    # Update batch size
+    payload['prompt']['232']['inputs']['batch_size'] = batch_size
+
+    # Update shift amount
+    payload['prompt']['67']['inputs']['shift'] = shift_amount
+
+    # Update guidance amount
+    payload['prompt']['26']['inputs']['guidance'] = guidance_amount
+
+    # Update steps amount
+    payload['prompt']['17']['inputs']['steps'] = steps_amount
+
+    # Update denoise amount
+    payload['prompt']['17']['inputs']['denoise'] = denoise_amount
+
+    # Update frame rate
+    payload['prompt']['82']['inputs']['frame_rate'] = frame_rate
+
+    # Update prompt text
+    payload['prompt']['44']['inputs']['text'] = prompt
+
+    return payload
+
+def post_prompt(url, payload):
+    """
+    Posts the prompt payload to the ComfyUI API.
+
+    Args:
+        url (str): ComfyUI API endpoint URL.
+        payload (dict): Prompt payload as a dictionary.
+
+    Returns:
+        str: Prompt ID if successful, None otherwise.
+    """
+    try:
+        response = requests.post(url, json=payload)
+        response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
+        return response.json().get('prompt_id')
+    except requests.exceptions.RequestException as e:
+        print(f"Error posting prompt: {e}")
+        return None
+
+def get_queue_progress(url):
+    """
+    Retrieves the queue progress from the ComfyUI API.
+
+    Args:
+        url (str): ComfyUI API queue endpoint URL.
+
+    Returns:
+        tuple: (queue_pending, queue_running) where queue_pending and queue_running are lists,
+               or (None, None) if an error occurred.
+    """
+    try:
+        response = requests.get(url)
+        response.raise_for_status()
+        queue_data = response.json()
+        queue_pending = queue_data.get('queue_pending', [])
+        queue_running = queue_data.get('queue_running', [])
+        return queue_pending, queue_running
+    except requests.exceptions.RequestException as e:
+        print(f"Error getting queue progress: {e}")
+        return None, None
+
+def get_history(url, prompt_id):
+    """
+    Retrieves the history from the ComfyUI API.
+
+    Args:
+        url (str): ComfyUI API history endpoint URL.
+        prompt_id (str): Prompt ID to look for in the history.
+
+    Returns:
+        tuple: (status, video_path) where status is a string ("success" or "failure") and video_path is the path to the video,
+               or (None, None) if an error occurred.
+    """
+    try:
+        response = requests.get(url)
+        response.raise_for_status()
+        history_data = response.json()
+
+        if history_data:
+            # Find the history item with the matching prompt_id
+            for history_item_id, history_item in history_data.items():
+                if history_item_id == prompt_id:
+                    status = history_item['status']['status_str']
+                    completed = history_item['status']['completed']
+
+                    if status == "success" and completed:
+                        try:
+                            for output in history_item['outputs'].values():
+                                for gif in output.get('gifs', []):
+                                    video_path = gif['fullpath']
+                                    return "success", video_path
+                        except (KeyError, TypeError):
+                            print("Video path not found in history data.")
+                            return "failure", None
+                    else:
+                        print("Video generation failed according to history data.")
+                        return "failure", None
+            print("Prompt ID not found in history data.")
+            return "failure", None
+        else:
+            print("No history data found.")
+            return "failure", None
+
+    except requests.exceptions.RequestException as e:
+        print(f"Error getting history: {e}")
+        return None, None
+
+def extract_video_number(filename):
+    """
+    Extracts the video number from the filename using a regular expression.
+    """
+    match = re.search(r'_(\d+)\.mp4', filename)
+    if match:
+        return match.group(1)
+    return None
+
+def main():
+    parser = argparse.ArgumentParser(description="Run ComfyUI with specified parameters.")
+    parser.add_argument("--width", type=int, default=320, help="Width of the generated image.")
+    parser.add_argument("--height", type=int, default=480, help="Height of the generated image.")
+    parser.add_argument("--length", type=int, default=45, help="Length (number of frames) of the video.")
+    parser.add_argument("--batchSize", type=int, default=1, help="Batch size for video generation.")
+    parser.add_argument("--shiftAmount", type=float, default=9.0, help="Shift amount for ModelSamplingSD3.")
+    parser.add_argument("--guidanceAmount", type=float, default=7.5, help="Guidance amount for FluxGuidance.")
+    parser.add_argument("--stepsAmount", type=int, default=12, help="Steps amount for BasicScheduler.")
+    parser.add_argument("--denoiseAmount", type=float, default=1.0, help="Denoise amount for BasicScheduler.")
+    parser.add_argument("--frameRate", type=int, default=24, help="Frame rate for video combine.")
+    parser.add_argument("--prompt", type=str, required=True, help="Text prompt for video generation.")
+    parser.add_argument("--upscale", type=bool, default=False, help="Whether to use upscale or not.")
+    parser.add_argument("--comfyui_url", type=str, default="http://localhost:8188", help="ComfyUI base URL.")
+
+    args = parser.parse_args()
+
+    API_URL = f"{args.comfyui_url}/api/prompt"
+    QUEUE_URL = f"{args.comfyui_url}/api/queue"
+    HISTORY_URL = f"{args.comfyui_url}/api/history?max_items=1"
+
+    if args.upscale:
+        prompt_file = "prompt-payload.json"
+        filename_prefix = "Hunyuan_upscaled"
+    else:
+        prompt_file = "prompt-payload-no-upscale.json"
+        filename_prefix = "Hunyuan_raw"
+
+    # Check if ComfyUI is running
+    print("Checking if ComfyUI is running...")
+    while not check_comfyui_started(args.comfyui_url):
+        print("ComfyUI not yet running.  Waiting...")
+        time.sleep(5)
+    print("ComfyUI is running.")
+
+    # Load and update the prompt payload
+    print("Loading and updating prompt payload...")
+    prompt_file_path = os.path.join(os.path.dirname(__file__), prompt_file)
+    try:
+        prompt_payload = load_and_update_prompt(prompt_file_path, args.width, args.height, args.length, args.batchSize, args.shiftAmount, args.guidanceAmount, args.stepsAmount, args.denoiseAmount, args.frameRate, args.prompt)
+    except Exception as e:
+        print(f"Error loading and updating prompt: {e}")
+        return
+
+    # Post the prompt
+    print("Posting prompt to ComfyUI...")
+    prompt_id = post_prompt(API_URL, prompt_payload)
+
+    if prompt_id:
+        print(f"Prompt submitted successfully with ID: {prompt_id}")
+
+        # Monitor the queue until it's empty
+        print("Monitoring queue...")
+        while True:
+            queue_pending, queue_running = get_queue_progress(QUEUE_URL)
+            if queue_pending is None or queue_running is None:
+                print("Failed to get queue progress.  Aborting.")
+                break
+
+            if not queue_pending and not queue_running:
+                print("Queue is empty.")
+                break
+            else:
+                print(f"Queue pending: {len(queue_pending)}, running: {len(queue_running)}.  Waiting...")
+                time.sleep(5)
+
+        # Get the history
+        print("Getting history...")
+        status, video_path = get_history(HISTORY_URL, prompt_id)
+
+        if status == "success":
+            print(f"Video generation complete! Video path: {video_path}")
+            video_number = extract_video_number(video_path)
+            if video_number:
+                output_filename = f"{filename_prefix}_{video_number}.mp4"
+            else:
+                output_filename = f"{filename_prefix}_latest.mp4"
+            print(output_filename)  # Print the output filename
+        else:
+            print("Video generation failed.")
+    else:
+        print("Failed to submit prompt.")
+
+if __name__ == "__main__":
+    main()