Files
lumalabs-ComfyUI-LumaAI-API/py/lumaai_api_node.py
T
2024-09-19 17:53:55 -03:00

402 lines
12 KiB
Python

import configparser
import os
import requests
import time
from lumaai import LumaAI
import folder_paths
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
config_path = os.path.join(parent_dir, "config.ini")
config = configparser.ConfigParser()
config.read(config_path)
try:
luma_api_key = config['API']['LUMAAI_API_KEY']
if luma_api_key != "":
os.environ["LUMAAI_API_KEY"] = luma_api_key
else:
print("Warning: LUMAAI_API_KEY is empty in config.ini")
except KeyError:
print("Error: LUMAAI_API_KEY not found in config.ini")
def download_file(url, file_name):
response = requests.get(url, stream=True)
with open(file_name, "wb") as file:
file.write(response.content)
print(f"File downloaded as {file_name}")
def parse_filename(filename):
# Remove file extension if present
filename = os.path.splitext(filename)[0]
# Check if the filename ends with a directory separator
if filename.endswith(os.path.sep):
# If it does, we'll use the generation_id as the filename later
directory = filename
filename = ""
else:
# Otherwise, split the path into directory and filename
directory, filename = os.path.split(filename)
# Create directories if they don't exist
output_path = folder_paths.get_output_directory()
full_directory = os.path.join(output_path, directory)
if directory and not os.path.exists(full_directory):
os.makedirs(full_directory)
return directory, filename
class LumaAIClient:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"api_key": (
"STRING",
{
"default": "",
},
)
},
}
RETURN_TYPES = ("LUMACLIENT",)
RETURN_NAMES = ("client",)
FUNCTION = "run"
CATEGORY = "LumaAI"
def run(self, api_key):
"""
Create a LumaAI client with the provided API key.
"""
api_key = api_key if api_key != "" else os.environ.get("LUMAAI_API_KEY", "")
if api_key == "":
raise ValueError("API Key is required")
client = LumaAI(
auth_token=api_key,
)
return (client,)
class Text2Video:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"client": ("LUMACLIENT", {"forceInput": True}),
"prompt": ("STRING", {"multiline": True, "default": ""}),
"loop": ("BOOLEAN", {"default": False}),
"aspect_ratio": (["9:16", "3:4", "1:1", "4:3", "16:9", "21:9"],),
"save": ("BOOLEAN", {"default": True}),
},
"optional": {"filename": ("STRING", {"default": ""})},
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("video_url", "generation_id")
OUTPUT_NODE = True
FUNCTION = "run"
CATEGORY = "LumaAI"
def run(self, client, prompt, loop, aspect_ratio, save, filename):
"""
Generate a video from a text prompt.
"""
if prompt == "":
raise ValueError("Prompt is required")
generation = client.generations.create(
prompt=prompt, loop=loop, aspect_ratio=aspect_ratio
)
generation_id = generation.id
completed = False
while not completed:
generation = client.generations.get(id=generation_id)
if generation.state == "completed":
completed = True
elif generation.state == "failed":
raise ValueError(f"Generation failed: {generation.failure_reason}")
time.sleep(3)
video_url = generation.assets.video
if save:
directory, filename = parse_filename(filename)
if filename == "":
filename = generation_id
download_file(video_url, os.path.join(self.output_dir, directory, filename + ".mp4"))
return {
"ui": {"text": [generation_id]},
"result": (
video_url,
generation_id,
),
}
class Image2Video:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"client": ("LUMACLIENT", {"forceInput": True}),
"prompt": ("STRING", {"multiline": True, "default": ""}),
"loop": ("BOOLEAN", {"default": False}),
"save": ("BOOLEAN", {"default": True}),
},
"optional": {
"init_image_url": ("STRING", {"default": "", "forceInput": True}),
"final_image_url": ("STRING", {"default": "", "forceInput": True}),
"filename": ("STRING", {"default": ""}),
},
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("video_url", "generation_id")
FUNCTION = "run"
CATEGORY = "LumaAI"
def run(
self,
client,
prompt,
loop,
save,
init_image_url="",
final_image_url="",
filename="",
):
"""
Generate a video from an image prompt.
"""
if init_image_url == "" and final_image_url == "":
raise ValueError("At least one image URL is required")
keyframes = {}
if init_image_url != "":
keyframes["frame0"] = {"type": "image", "url": init_image_url}
if final_image_url != "":
keyframes["frame1"] = {"type": "image", "url": final_image_url}
generation = client.generations.create(
prompt=prompt, loop=loop, keyframes=keyframes
)
generation_id = generation.id
completed = False
while not completed:
generation = client.generations.get(id=generation_id)
if generation.state == "completed":
completed = True
elif generation.state == "failed":
raise ValueError(f"Generation failed: {generation.failure_reason}")
time.sleep(3)
video_url = generation.assets.video
if save:
directory, filename = parse_filename(filename)
if filename == "":
filename = generation_id
download_file(video_url, os.path.join(self.output_dir, directory, filename + ".mp4"))
return {
"ui": {"text": [generation_id]},
"result": (
video_url,
generation_id,
),
}
class InterpolateGenerations:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"client": ("LUMACLIENT", {"forceInput": True}),
"prompt": ("STRING", {"multiline": True, "default": ""}),
"save": ("BOOLEAN", {"default": True}),
"generation_id_1": ("STRING", {"default": "", "forceInput": True}),
"generation_id_2": ("STRING", {"default": "", "forceInput": True}),
},
"optional": {"filename": ("STRING", {"default": ""})},
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("video_url", "generation_id")
FUNCTION = "run"
CATEGORY = "LumaAI"
def run(
self,
client,
prompt,
loop,
aspect_ratio,
save,
generation_id_1,
generation_id_2,
filename="",
):
"""
Generate a video by interpolating between two existing generations.
"""
if not generation_id_1 or not generation_id_2:
raise ValueError("Both generation IDs are required")
generation = client.generations.create(
prompt=prompt,
loop=loop,
aspect_ratio=aspect_ratio,
keyframes={
"frame0": {"type": "generation", "id": generation_id_1},
"frame1": {"type": "generation", "id": generation_id_2},
},
)
generation_id = generation.id
completed = False
while not completed:
generation = client.generations.get(id=generation_id)
if generation.state == "completed":
completed = True
elif generation.state == "failed":
raise ValueError(f"Generation failed: {generation.failure_reason}")
time.sleep(3)
video_url = generation.assets.video
if save:
directory, filename = parse_filename(filename)
if filename == "":
filename = generation_id
download_file(video_url, os.path.join(self.output_dir, directory, filename + ".mp4"))
return {
"ui": {"text": [generation_id]},
"result": (
video_url,
generation_id,
),
}
class ExtendGeneration:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"client": ("LUMACLIENT", {"forceInput": True}),
"prompt": ("STRING", {"multiline": True, "default": ""}),
"save": ("BOOLEAN", {"default": True}),
},
"optional": {
"init_image_url": ("STRING", {"default": "", "forceInput": True}),
"final_image_url": ("STRING", {"default": "", "forceInput": True}),
"init_generation_id": ("STRING", {"default": "", "forceInput": True}),
"final_generation_id": ("STRING", {"default": "", "forceInput": True}),
"filename": ("STRING", {"default": ""}),
},
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("video_url", "generation_id")
FUNCTION = "run"
CATEGORY = "LumaAI"
def run(
self,
client,
prompt,
save,
init_image_url="",
final_image_url="",
init_generation_id="",
final_generation_id="",
filename="",
):
"""
Generate a video by extending from an image to an existing generation.
"""
if not init_generation_id and not final_generation_id:
raise ValueError("You must provide at least one generation id")
if init_image_url and init_generation_id:
raise ValueError(
"You cannot provide both an init image and a init generation"
)
if final_image_url and final_generation_id:
raise ValueError(
"You cannot provide both a final image and a final generation"
)
keyframes = {}
if init_image_url != "":
keyframes["frame0"] = {"type": "image", "url": init_image_url}
if final_image_url != "":
keyframes["frame1"] = {"type": "image", "url": final_image_url}
if init_generation_id != "":
keyframes["frame0"] = {"type": "generation", "id": init_generation_id}
if final_generation_id != "":
keyframes["frame1"] = {"type": "generation", "id": final_generation_id}
generation = client.generations.create(prompt=prompt, keyframes=keyframes)
new_generation_id = generation.id
completed = False
while not completed:
generation = client.generations.get(id=new_generation_id)
if generation.state == "completed":
completed = True
elif generation.state == "failed":
raise ValueError(f"Generation failed: {generation.failure_reason}")
time.sleep(3)
video_url = generation.assets.video
if save:
directory, filename = parse_filename(filename)
if filename == "":
filename = new_generation_id
download_file(video_url, os.path.join(self.output_dir, directory, filename + ".mp4"))
return {
"ui": {"text": [new_generation_id]},
"result": (
video_url,
new_generation_id,
),
}
class PreviewVideo:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"video_url": ("STRING", {"forceInput": True}),
}
}
OUTPUT_NODE = True
FUNCTION = "run"
CATEGORY = "LumaAI"
RETURN_TYPES = ()
def run(self, video_url):
return {"ui": {"video_url": [video_url]}}