nodes/ru4ls-ComfyUI_Wan
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
69ddf0d20356d49feceacbae60d51941b2630dea
ru4ls-ComfyUI_Wan/generators/ii2v.py
T
ru4ls b22ecfb7da Enhance region selection support for Wan API 
- Updated .env.template to include separate API keys for international and mainland China endpoints.
- Modified base.py to load environment variables from multiple .env file locations and support region-specific API endpoints.
- Refactored API endpoint definitions to use a centralized dictionary for international and mainland China endpoints.
- Updated generator classes (i2v, ii2v, t2i, t2v, VACE) to accept region as a parameter and utilize the appropriate API endpoints and keys.
- Implemented region selection in the video processing classes, ensuring correct API requests based on the selected region.
2025-09-13 13:49:16 +07:00

290 lines
12 KiB
Python
Raw Blame History

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 WanII2VGenerator(WanAPIBase):
"""Node for image-to-video generation using first and last frames with Wan model"""
# Define available Wan ii2v models
MODEL_OPTIONS = [
"wan2.1-kf2v-plus" # Professional Edition
]
# Define allowed resolutions for Wan ii2v models
RESOLUTION_OPTIONS = [
"720P"
]
# 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-kf2v-plus"
}),
"first_frame_url": ("STRING", {
"default": "https://example.com/first_frame.png"
}),
"last_frame_url": ("STRING", {
"default": "https://example.com/last_frame.png"
}),
"prompt": ("STRING", {
"multiline": True,
"default": "A black kitten looks up at the sky curiously, the camera gradually rises from eye level, and finally shoots from a top-down angle to capture the kitten's curious eyes."
}),
"region": (cls.REGION_OPTIONS, {
"default": "international"
})
},
"optional": {
"negative_prompt": ("STRING", {
"multiline": True,
"default": ""
}),
"resolution": (cls.RESOLUTION_OPTIONS, {
"default": "720P"
}),
"prompt_extend": ("BOOLEAN", {
"default": True
}),
"watermark": ("BOOLEAN", {
"default": False
}),
"seed": ("INT", {
"default": 0,
"min": 0,
"max": 2147483647
}),
"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"
def generate(self, model, first_frame_url, last_frame_url, prompt, region, negative_prompt="",
resolution="720P", prompt_extend=True, watermark=False, seed=0, 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["ii2v_post"]
# Prepare API payload for image-to-video generation with first and last frames
payload = {
"model": model,
"input": {
"first_frame_url": first_frame_url,
"last_frame_url": last_frame_url,
"prompt": prompt
},
"parameters": {
"resolution": resolution,
"prompt_extend": prompt_extend,
"watermark": watermark
}
}
# Add optional parameters if they have non-default values
if negative_prompt:
payload["input"]["negative_prompt"] = negative_prompt
if seed > 0:
payload["parameters"]["seed"] = seed
# 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}")
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_ii2v_{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)")
Reference in New Issue View Git Blame Copy Permalink