29 lines
719 B
Python
29 lines
719 B
Python
import os
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
class CozyPoseReferenceBodyCustomNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"angle": ("INT", {"default": 0, "min": 0, "max": 315, "step": 45}),
|
|
"size": (["512", "1024"], {}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("pose",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "CozyMantis"
|
|
|
|
def run(self, angle, size):
|
|
pose_path = os.path.join(os.path.dirname(__file__), "poses", "body", size, f"{angle}.png")
|
|
image = Image.open(pose_path)
|
|
image = image.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
|
|
return (image,)
|