290 lines
12 KiB
Python
290 lines
12 KiB
Python
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 WanT2VGenerator(WanAPIBase):
|
|
"""Node for text-to-video generation using Wan model"""
|
|
|
|
# Define available Wan t2v models
|
|
MODEL_OPTIONS = [
|
|
"wan2.2-t2v-plus" # Professional Edition
|
|
]
|
|
|
|
# Define allowed resolutions for Wan t2v models
|
|
RESOLUTION_OPTIONS = [
|
|
"480P",
|
|
"1080P"
|
|
]
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
# Use the centralized API endpoint from the base class
|
|
# To use Mainland China region, modify API_ENDPOINT_POST_VIDEO in core/base.py
|
|
self.api_url = self.API_ENDPOINT_POST_VIDEO
|
|
|
|
@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.2-t2v-plus"
|
|
}),
|
|
"prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "A kitten running in the moonlight"
|
|
})
|
|
},
|
|
"optional": {
|
|
"negative_prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": ""
|
|
}),
|
|
"resolution": (cls.RESOLUTION_OPTIONS, {
|
|
"default": "1080P"
|
|
}),
|
|
"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, prompt, negative_prompt="", resolution="1080P",
|
|
prompt_extend=True, watermark=False, seed=0, output_dir="./videos"):
|
|
# Check API key
|
|
self.check_api_key()
|
|
|
|
# Prepare API payload for text-to-video generation
|
|
payload = {
|
|
"model": model,
|
|
"input": {
|
|
"prompt": prompt
|
|
},
|
|
"parameters": {
|
|
"prompt_extend": prompt_extend,
|
|
"watermark": watermark
|
|
}
|
|
}
|
|
|
|
# Add resolution parameter based on selection
|
|
# Convert resolution tier to specific size values
|
|
resolution_sizes = {
|
|
"480P": {
|
|
"16:9": "832*480",
|
|
"9:16": "480*832",
|
|
"1:1": "624*624"
|
|
},
|
|
"1080P": {
|
|
"16:9": "1920*1080",
|
|
"9:16": "1080*1920",
|
|
"1:1": "1440*1440",
|
|
"4:3": "1632*1248",
|
|
"3:4": "1248*1632"
|
|
}
|
|
}
|
|
|
|
# Default to 16:9 aspect ratio
|
|
if resolution in resolution_sizes:
|
|
payload["parameters"]["size"] = resolution_sizes[resolution]["16:9"]
|
|
|
|
# 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 {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
"X-DashScope-Async": "enable" # Wan requires async processing
|
|
}
|
|
|
|
try:
|
|
# Make API request
|
|
print(f"Making API request to {self.api_url}")
|
|
response = requests.post(self.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)
|
|
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"):
|
|
"""Poll for task result until completion and download video"""
|
|
import time
|
|
|
|
# URL for querying task results
|
|
# To use Mainland China region, modify API_ENDPOINT_GET in core/base.py
|
|
query_url = self.API_ENDPOINT_GET.format(task_id=task_id)
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {self.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_t2v_{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)") |