""" Wan VACE Multi-Image Reference Node for ComfyUI """ import os import json import requests from PIL import Image import numpy as np import torch import io import base64 from dotenv import load_dotenv import sys import pathlib from datetime import datetime # Import the base class and COMFYUI_AVAILABLE flag from ..core.base import WanAPIBase, COMFYUI_AVAILABLE # Try to import folder_paths if available try: import folder_paths except ImportError: pass class WanVACEImageReference(WanAPIBase): """Node for multi-image reference using Wan VACE model""" # Define available Wan VACE models MODEL_OPTIONS = [ "wan2.1-vace-plus" # Professional Edition ] # Define video resolutions RESOLUTION_OPTIONS = [ "1280*720", # 16:9 aspect ratio (default) "720*1280", # 9:16 aspect ratio "960*960", # 1:1 aspect ratio "832*1088", # 3:4 aspect ratio "1088*832" # 4:3 aspect ratio ] # Define region options REGION_OPTIONS = [ "international", "mainland_china" ] def __init__(self): super().__init__() @classmethod def INPUT_TYPES(cls): # Define output directory options if COMFYUI_AVAILABLE: # Use ComfyUI's output directory with browseable option output_dir_options = { "default": "./videos", "tooltip": "Directory where the generated video will be saved. Browse to select a custom directory." } else: # Fallback to string input output_dir_options = { "default": "./videos", "multiline": False } return { "required": { "model": (cls.MODEL_OPTIONS, { "default": "wan2.1-vace-plus" }), "prompt": ("STRING", { "multiline": True, "default": "Generate a video based on the provided reference images" }), "ref_images_url": ("STRING", { "multiline": True, "default": "", "tooltip": "Newline-separated URLs for reference images" }), "region": (cls.REGION_OPTIONS, { "default": "international" }) }, "optional": { "obj_or_bg": ("STRING", { "multiline": True, "default": "", "tooltip": "Newline-separated values (obj/bg) corresponding to ref_images_url" }), "size": (cls.RESOLUTION_OPTIONS, { "default": "1280*720" }), "seed": ("INT", { "default": 0, "min": 0, "max": 2147483647 }), "prompt_extend": ("BOOLEAN", { "default": False }), "watermark": ("BOOLEAN", { "default": False }), "output_dir": ("STRING", output_dir_options) } } RETURN_TYPES = ("STRING", "STRING") # Returns path to downloaded video file and video URL RETURN_NAMES = ("video_file_path", "video_url") FUNCTION = "generate" CATEGORY = "Ru4ls/Wan/VACE" def generate(self, model, prompt, ref_images_url, region, obj_or_bg="", size="1280*720", seed=0, prompt_extend=False, watermark=False, output_dir="./videos"): # Check API key based on region api_key = self.check_api_key(region) # Get the appropriate API endpoints based on region endpoints = self.get_api_endpoints(region) api_url = endpoints["video_post"] # Validate required inputs if not ref_images_url or not ref_images_url.strip(): raise ValueError("Reference images URLs are required for image reference function") # Handle ref_images_url as a list ref_images_list = [url.strip() for url in ref_images_url.split('\n') if url.strip()] if not ref_images_list: raise ValueError("At least one valid reference image URL is required") # Prepare API payload payload = { "model": model, "input": { "function": "image_reference", "prompt": prompt, "ref_images_url": ref_images_list }, "parameters": { "size": size, "prompt_extend": prompt_extend, "watermark": watermark } } # Add seed if provided if seed > 0: payload["parameters"]["seed"] = seed # Handle obj_or_bg as a list if obj_or_bg: obj_or_bg_list = [item.strip() for item in obj_or_bg.split('\n') if item.strip()] if obj_or_bg_list: # Validate that obj_or_bg_list has the same length as ref_images_list if len(obj_or_bg_list) != len(ref_images_list): raise ValueError("obj_or_bg list must have the same length as ref_images_url list") payload["parameters"]["obj_or_bg"] = obj_or_bg_list else: # If obj_or_bg is not provided, automatically generate it: # - For single image: default to ["obj"] # - For multiple images: assume last image is background, rest are entities if len(ref_images_list) == 1: payload["parameters"]["obj_or_bg"] = ["obj"] else: # For multiple images, assume last is bg, rest are obj obj_or_bg_auto = ["obj"] * (len(ref_images_list) - 1) + ["bg"] payload["parameters"]["obj_or_bg"] = obj_or_bg_auto # Set headers according to DashScope documentation headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", "X-DashScope-Async": "enable" # Wan requires async processing } try: # Make API request print(f"Making API request to {api_url}") print(f"Payload: {json.dumps(payload, indent=2)}") response = requests.post(api_url, headers=headers, json=payload) print(f"Response status code: {response.status_code}") if hasattr(response, 'text'): print(f"Response text: {response.text[:500]}...") # Print first 500 chars response.raise_for_status() # Parse response to get task_id result = response.json() print(f"API response received: {json.dumps(result, indent=2)[:200]}...") # Print first 200 chars # Check if this is a task creation response if "output" in result and "task_id" in result["output"]: task_id = result["output"]["task_id"] task_status = result["output"]["task_status"] print(f"Task created with ID: {task_id}, status: {task_status}") # Now we need to poll for the result task_result = self.poll_task_result(task_id, output_dir, region) return task_result # Return both path to downloaded video file and video URL else: raise ValueError(f"Unexpected API response format: {result}") except requests.exceptions.RequestException as e: # More detailed error handling if hasattr(e, 'response') and e.response is not None: status_code = e.response.status_code response_text = e.response.text print(f"API request failed with status {status_code}: {response_text}") if status_code == 401: raise RuntimeError(f"API request failed: 401 Unauthorized. " f"This usually means your API key is invalid or not properly configured. " f"Error details: {response_text}") elif status_code == 403: raise RuntimeError(f"API request failed: 403 Forbidden. " f"This usually means your API key is valid but you don't have access to this model. " f"Error details: {response_text}") elif status_code == 400: raise RuntimeError(f"API request failed: 400 Bad Request. " f"This usually means there's an issue with the request format. " f"Error details: {response_text}") else: raise RuntimeError(f"API request failed: {status_code} {e.response.reason}. Response: {response_text}") else: raise RuntimeError(f"API request failed: {str(e)}") except Exception as e: raise RuntimeError(f"Failed to process API response: {str(e)}") def poll_task_result(self, task_id, output_dir="./videos", region="international"): """Poll for task result until completion and download video""" import time # Get the appropriate API endpoints based on region endpoints = self.get_api_endpoints(region) query_url = endpoints["get"].format(task_id=task_id) # Check API key based on region api_key = self.check_api_key(region) headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json" } max_attempts = 60 # Maximum polling attempts (may take longer for video) attempt = 0 while attempt < max_attempts: try: print(f"Polling task {task_id}, attempt {attempt + 1}/{max_attempts}") response = requests.get(query_url, headers=headers) response.raise_for_status() result = response.json() task_status = result["output"]["task_status"] print(f"Task status: {task_status}") if task_status == "SUCCEEDED": # Task completed successfully if "video_url" in result["output"]: video_url = result["output"]["video_url"] # Download the video video_response = requests.get(video_url) video_response.raise_for_status() # Create a unique filename for the video timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") video_filename = f"wan_vace_image_reference_{timestamp}.mp4" # Handle output directory based on ComfyUI availability if COMFYUI_AVAILABLE and not output_dir.startswith(("./", "/")): # Use ComfyUI's output directory structure if output_dir.endswith("/"): output_dir = output_dir[:-1] full_output_folder = folder_paths.get_output_directory() output_path = os.path.join(full_output_folder, output_dir) else: # Resolve output directory path (existing logic) if output_dir.startswith("./"): # Relative to the node directory output_path = os.path.join(os.path.dirname(__file__), output_dir[2:]) else: output_path = output_dir # Create output directory if it doesn't exist os.makedirs(output_path, exist_ok=True) # Save video to file video_path = os.path.join(output_path, video_filename) with open(video_path, "wb") as f: f.write(video_response.content) print(f"Video downloaded and saved to: {video_path}") # Return path relative to ComfyUI output directory if using ComfyUI if COMFYUI_AVAILABLE and not output_dir.startswith(("./", "/")): return_path = os.path.join(output_dir, video_filename) if output_dir != "./videos" else video_filename else: return_path = video_path # Return full path # Return both the file path and the video URL return (return_path, video_url) else: raise ValueError(f"Unexpected API response format: {result}") elif task_status == "FAILED": # Task failed error_code = result["output"].get("code", "Unknown") error_message = result["output"].get("message", "Unknown error") raise RuntimeError(f"Task failed with code: {error_code}, message: {error_message}") elif task_status in ["PENDING", "RUNNING"]: # Task still in progress, wait and retry time.sleep(10) # Wait 10 seconds before retrying (video generation may take longer) attempt += 1 continue else: raise ValueError(f"Unexpected task status: {task_status}") except requests.exceptions.RequestException as e: raise RuntimeError(f"Failed to query task status: {str(e)}") # If we've reached here, we've exceeded max attempts raise RuntimeError(f"Task did not complete within the expected time ({max_attempts} attempts)")