Update run_comfyui.py

run_comfyui.py

@@ -1,23 +1,9 @@
-import os
-import subprocess
-import boto3
-from datetime import datetime
-from typing import Optional
-from src.models import GenerateVideoRequest
-import modal
-import time
 import requests
-
-# Define Modal stub and AWS secrets
-stub = modal.App("cui-hunyuan-video")
-aws_secret = modal.Secret.from_name("aws-s3-video-bucket")
-
-# Define Docker image
-dockerfile_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../models/CUI-Hunyuan/resources/Dockerfile"))
-image = modal.Image.from_dockerfile(dockerfile_path).pip_install("boto3")
-
-# AWS S3 Configuration
-S3_BUCKET = "flickify-generated-videos"
+import json
+import time
+import argparse
+import os
+import re
 
 def check_comfyui_started(url):
     """
@@ -25,141 +11,245 @@ def check_comfyui_started(url):
     """
     try:
         response = requests.get(url)
-        print(f"COMFYUI STATUS CHECK: ComfyUI response: {response.status_code}")
         return response.status_code == 200
     except requests.ConnectionError:
         return False
 
-@stub.function(
-    image=image,
-    gpu="a10g",
-    timeout=1800,
-    secrets=[aws_secret]
-)
-def run_cui_hunyuan_video(prompt: str, width: Optional[int] = None, height: Optional[int] = None, length: Optional[int] = None, batch_size: Optional[int] = None, shift_amount: Optional[float] = None, guidance_amount: Optional[float] = None, steps_amount: Optional[int] = None, denoise_amount: Optional[float] = None, frame_rate: Optional[int] = None, upscale: Optional[bool] = None):
-    
+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.
     """
-    Runs the CUI-Hunyuan Video model, uploads to S3, and returns a pre-signed URL.
+    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):
     """
-    print('Starting ComfyUI...')
-    # Start ComfyUI as a background process
-    comfyui_process = subprocess.Popen(
-        ["python", "main.py", "--port", "8188", "--listen"],
-        cwd="/workspace/ComfyUI"
-    )
+    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:
-        # Check if ComfyUI is running
-        comfyui_url = "http://localhost:8188"
-        for _ in range(20):
-            if check_comfyui_started(comfyui_url):
-                print("ComfyUI is running.")
-                break
-            time.sleep(5)
+        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("Failed to start ComfyUI.")
-            raise RuntimeError("Failed to start ComfyUI.")
-
-        print('Running CUI-Hunyuan Video model...')
-        # Construct the command to run the video generation script
-        command = [
-            "python3", "run_comfyui.py",
-            "--prompt", prompt
-        ]
-
-        if width is not None:
-            command.extend(["--width", str(width)])
-        if height is not None:
-            command.extend(["--height", str(height)])
-        if length is not None:
-            command.extend(["--length", str(length)])
-        if batch_size is not None:
-            command.extend(["--batchSize", str(batch_size)])
-        if shift_amount is not None:
-            command.extend(["--shiftAmount", str(shift_amount)])
-        if guidance_amount is not None:
-            command.extend(["--guidanceAmount", str(guidance_amount)])
-        if steps_amount is not None:
-            command.extend(["--stepsAmount", str(steps_amount)])
-        if denoise_amount is not None:
-            command.extend(["--denoiseAmount", str(denoise_amount)])
-        if frame_rate is not None:
-            command.extend(["--frameRate", str(frame_rate)])
-        if upscale is not None:
-            command.extend(["--upscale", str(upscale)])
-
-        # Run the model
-        process = subprocess.run(command, shell=False, capture_output=True, text=True)
-
-        if process.returncode != 0:
-            raise RuntimeError(f"Error running model: {process.stderr}")
-
-        # Capture the output filename from the script
-        output_lines = process.stdout.strip().split('\n')
-        output_filename = output_lines[-1]  # Assuming the last line is the filename
-
-        # Check if the output filename is valid
-        if not output_filename or "Error" in output_filename:
-            raise RuntimeError(f"Error in output: {process.stdout}")
-
-        output_file = os.path.join("/workspace/ComfyUI/output", output_filename)
-
-        # Check if the output file exists
-        if not os.path.exists(output_file):
-            raise RuntimeError(f"Output file not found: {output_file}")
-
-        # Upload to S3
-        try:
-            s3 = boto3.client("s3")
-            timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-            new_filename = f"{timestamp}_{output_filename}"
-            with open(output_file, 'rb') as file_obj:
-                s3.upload_fileobj(
-                    file_obj,
-                    S3_BUCKET,
-                    f"videos/{new_filename}"
-                )
-
-            # Generate a pre-signed URL that expires in 1 hour (3600 seconds)
-            video_url = s3.generate_presigned_url(
-                'get_object',
-                Params={'Bucket': S3_BUCKET, 'Key': f"videos/{new_filename}"},
-                ExpiresIn=3600
-            )
-
-            return {
-                "message": "Video generation successful",
-                "video_url": video_url
-            }
-
-        except Exception as e:
-            raise RuntimeError(f"Failed to upload video to S3: {str(e)}")
-
-    finally:
-        comfyui_process.terminate()
-
-async def generate_video_cui_hunyuan(request: GenerateVideoRequest) -> dict:
+            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):
     """
-    Handles the video generation request and returns the result.
+    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:
-        async with stub.run():
-            print("Running CUI-Hunyuan Video model...")
-            # Call remote function without await, as it returns a Modal handle
-            result = run_cui_hunyuan_video.remote(
-                prompt=request.prompt,
-                width=request.width,
-                height=request.height,
-                length=request.length,
-                batch_size=request.batch_size,
-                shift_amount=request.shift_amount,
-                guidance_amount=request.guidance_amount,
-                steps_amount=request.steps_amount,
-                denoise_amount=request.denoise_amount,
-                frame_rate=request.frame_rate,
-                upscale=request.upscale
-            )
-            print("Result: ", await result)
-            return await result
+        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:
-        raise RuntimeError(f"Error generating video: {str(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()