Update run_comfyui.py

run_comfyui.py

@@ -1,255 +1,165 @@
-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()
+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"
+
+def check_comfyui_started(url):
+    """
+    Checks if ComfyUI is running at the given 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):
+    
+    """
+    Runs the CUI-Hunyuan Video model, uploads to S3, and returns a pre-signed URL.
+    """
+    print('Starting ComfyUI...')
+    # Start ComfyUI as a background process
+    comfyui_process = subprocess.Popen(
+        ["python", "main.py", "--port", "8188", "--listen"],
+        cwd="/workspace/ComfyUI"
+    )
+
+    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)
+        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:
+    """
+    Handles the video generation request and returns the result.
+    """
+    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
+    except Exception as e:
+        raise RuntimeError(f"Error generating video: {str(e)}")