122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
import os
|
||
import io
|
||
import json
|
||
import torch
|
||
from . import generator
|
||
from io import BytesIO
|
||
from PIL import Image
|
||
|
||
global luma_key_path
|
||
lunma_key_path = os.path.dirname(os.path.realpath(__file__))
|
||
|
||
def get_cookie_key():
|
||
try:
|
||
config_path = os.path.join(lunma_key_path, 'key.json')
|
||
with open(config_path, 'r') as f:
|
||
config = json.load(f)
|
||
api_key = config["cookie"]
|
||
except:
|
||
print("Error: Cookie is required. Please enter correct cookie in key.json")
|
||
return ""
|
||
return api_key
|
||
|
||
def create_folder(path):
|
||
try:
|
||
os.makedirs(path, exist_ok=True)
|
||
except Exception as e:
|
||
print(f"Fail to create the folder：{e}")
|
||
|
||
class LUMA_YoC_generator:
|
||
|
||
def __init__(self):
|
||
self.api_key = get_cookie_key()
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"prompt": ("STRING", {"default": "movie,cinematic", "multiline": True}),
|
||
"Enhance_prompt": ("BOOLEAN", {"default": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("luma_id",)
|
||
FUNCTION = "LUMA_generation"
|
||
CATEGORY = "YoC app/LumaAPP"
|
||
|
||
def LUMA_generation(self,image, prompt,Enhance_prompt):
|
||
tensor = image
|
||
if tensor.dim() == 4 and tensor.shape[0] == 1:
|
||
tensor = tensor.squeeze(0) # 去掉批次维度，形状变为[3, H, W]
|
||
|
||
# 如果通道在前，需要调整为[H, W, 3]以适应PIL或其他图像库
|
||
if tensor.shape[0] == 3:
|
||
tensor = tensor.permute(1, 2, 0)
|
||
# 确保tensor是浮点类型并且范围在[0, 1]，如果不是，您可能需要先进行归一化或缩放
|
||
tensor = tensor.clamp(0, 1) # 确保值在[0, 1]之间
|
||
|
||
# 判断图像是否为空白
|
||
if torch.std(tensor) < 0.01: # 可调整阈值以适应具体情况
|
||
print("Empty Image")
|
||
image_binary = None
|
||
else:
|
||
# 转换为NumPy数组
|
||
numpy_array = tensor.numpy()
|
||
buffer = BytesIO()
|
||
# 转换为PIL图像并保存
|
||
image = Image.fromarray((numpy_array * 255).astype('uint8')) # 转换为0-255范围
|
||
image.save(buffer, format='PNG') # 您可以选择其他格式，如JPEG
|
||
|
||
# 获取二进制数据
|
||
image_binary = buffer.getvalue()
|
||
|
||
res = generator.DreamMachine_generation(prompt=prompt, binary_image = image_binary,
|
||
Cookie=self.api_key, expand_prompt=Enhance_prompt)
|
||
res_id = res[0]['id']
|
||
print("Result id: %s"%res_id)
|
||
create_folder(os.path.join(os.getcwd(), "ComfyUI/output/luma video"))
|
||
with open(os.path.join(os.path.join(os.getcwd(), "ComfyUI/output/luma video"),"record_id.txt"), 'a') as file:
|
||
file.write(res_id + '\n')
|
||
|
||
return (res_id,)
|
||
|
||
class LUMA_YoC_result:
|
||
|
||
def __init__(self):
|
||
self.api_key = get_cookie_key()
|
||
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"text": ("STRING", {"default": "Enter your id", "multiline": False}),
|
||
"download":("BOOLEAN", {"default": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("video_link",)
|
||
FUNCTION = "LUMA_result"
|
||
CATEGORY = "YoC app/LumaAPP"
|
||
|
||
def LUMA_result(self,text,download):
|
||
video_folder = os.path.join(os.getcwd(), "ComfyUI/output/luma video")
|
||
create_folder(video_folder)
|
||
res_link = generator.get_signal_video(ids = text, Cookie = self.api_key,
|
||
download = download,output_name=os.path.join(video_folder,f"{text}.mp4"))
|
||
return (res_link,)
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"LUMA_API_YoC": LUMA_YoC_generator,
|
||
"LUMA_API_result_YoC": LUMA_YoC_result,
|
||
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"LUMA_API_YoC": "LUMA_generator @.@",
|
||
"LUMA_API_result_YoC": "LUMA_YoC_result @.@",
|
||
}
|