363 lines
13 KiB
Python
363 lines
13 KiB
Python
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
import folder_paths
|
|
|
|
from .fashn_vton import TryOnPipeline
|
|
from .fashn_vton.pose_keypoints import convert_pose_keypoints_to_dwpose
|
|
from .fashn_vton.preprocessing import FASHN_LABELS_TO_IDS
|
|
|
|
MODEL_LIST = ["fashn-ai/fashn-vton-1.5"]
|
|
MAX_FASHN_LABEL_ID = max(FASHN_LABELS_TO_IDS.values())
|
|
DEFAULT_CATEGORY_LABEL_ID = {
|
|
"tops": FASHN_LABELS_TO_IDS["top"],
|
|
"bottoms": FASHN_LABELS_TO_IDS["pants"],
|
|
"one-pieces": FASHN_LABELS_TO_IDS["dress"],
|
|
}
|
|
|
|
|
|
def _tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image:
|
|
"""Convert a ComfyUI IMAGE tensor (B,H,W,C in [0,1]) to PIL for pipeline input."""
|
|
img = image_tensor[0].detach().cpu().numpy()
|
|
img = np.clip(img, 0.0, 1.0)
|
|
|
|
if img.ndim == 2:
|
|
return Image.fromarray((img * 255.0).astype(np.uint8), mode="L")
|
|
|
|
channels = img.shape[2]
|
|
|
|
if channels == 1:
|
|
return Image.fromarray((img[..., 0] * 255.0).astype(np.uint8), mode="L")
|
|
|
|
if channels == 2:
|
|
return Image.fromarray((img[..., 0] * 255.0).astype(np.uint8), mode="L")
|
|
|
|
if channels > 3:
|
|
img = img[..., :3]
|
|
|
|
return Image.fromarray((img * 255.0).astype(np.uint8))
|
|
|
|
|
|
class FashnVtonLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": (MODEL_LIST, {"default": "fashn-ai/fashn-vton-1.5"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FASHN_VTON_PIPELINE",)
|
|
RETURN_NAMES = ("pipeline",)
|
|
FUNCTION = "load_pipeline"
|
|
CATEGORY = "FashnAI"
|
|
|
|
def load_pipeline(self, model):
|
|
weights_name = "fashn-vton"
|
|
base_weights_dir = os.path.join(folder_paths.models_dir, weights_name)
|
|
|
|
os.makedirs(base_weights_dir, exist_ok=True)
|
|
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
# Download TryOnModel
|
|
tryon_path = os.path.join(base_weights_dir, "model.safetensors")
|
|
if not os.path.exists(tryon_path):
|
|
print(f"FashnVTON: Downloading TryOnModel weights to {tryon_path}...")
|
|
hf_hub_download(
|
|
repo_id=model,
|
|
filename="model.safetensors",
|
|
local_dir=base_weights_dir,
|
|
)
|
|
|
|
# Download DWPose
|
|
dwpose_dir = os.path.join(base_weights_dir, "dwpose")
|
|
os.makedirs(dwpose_dir, exist_ok=True)
|
|
for filename in ["yolox_l.onnx", "dw-ll_ucoco_384.onnx"]:
|
|
if not os.path.exists(os.path.join(dwpose_dir, filename)):
|
|
print(f"FashnVTON: Downloading DWPose/{filename} to {dwpose_dir}...")
|
|
hf_hub_download(
|
|
repo_id="fashn-ai/DWPose",
|
|
filename=filename,
|
|
local_dir=dwpose_dir,
|
|
)
|
|
|
|
# Initialize Pipeline
|
|
print(f"FashnVTON: Loading pipeline from {base_weights_dir}...")
|
|
pipeline = TryOnPipeline(weights_dir=base_weights_dir)
|
|
|
|
return (pipeline,)
|
|
|
|
|
|
class FashnVtonInference:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("FASHN_VTON_PIPELINE",),
|
|
"person_image": ("IMAGE",),
|
|
"garment_image": ("IMAGE",),
|
|
"garment_photo_type": (
|
|
["model", "flat-lay"],
|
|
{
|
|
"default": "model",
|
|
"tooltip": "'model' if garment is worn by a person, 'flat-lay' for product shots on plain backgrounds",
|
|
},
|
|
),
|
|
"category": (
|
|
["tops", "bottoms", "one-pieces"],
|
|
{"default": "tops", "tooltip": "Garment category - 'tops', 'bottoms', or 'one-pieces'"},
|
|
),
|
|
"skip_cfg_last_n_steps": (
|
|
"INT",
|
|
{"default": 1, "tooltip": "Skip CFG for final N steps to prevent color saturation"},
|
|
),
|
|
"segmentation_free": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "If True, skip person clothing-agnostic masking. Set to False to use person segmentation for person mask control.",
|
|
},
|
|
),
|
|
"steps": (
|
|
"INT",
|
|
{
|
|
"default": 30,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1,
|
|
"tooltip": "Recommended: 20 (fast), 30 (balanced), 50 (quality)",
|
|
},
|
|
),
|
|
"cfg": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 10.0, "step": 0.1}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"keep_model_loaded": ("BOOLEAN", {"default": True}),
|
|
"pose_source": (
|
|
["auto", "internal_dwpose", "external_pose_keypoints"],
|
|
{
|
|
"default": "auto",
|
|
"tooltip": "auto: use external keypoints when connected, otherwise internal DWPose fallback.",
|
|
},
|
|
),
|
|
"parser_backend": (
|
|
["fashn_human_parser", "external_fashn_labelmap"],
|
|
{
|
|
"default": "fashn_human_parser",
|
|
"tooltip": "Use external_fashn_labelmap only when providing segmentation images encoded with FASHN label IDs (0-17).",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"person_pose_keypoints": (
|
|
"FASHN_DWPOSE_KEYPOINTS",
|
|
{"tooltip": "Optional external keypoints for person pose. Used when pose_source is auto/external_pose_keypoints."},
|
|
),
|
|
"garment_pose_keypoints": (
|
|
"FASHN_DWPOSE_KEYPOINTS",
|
|
{"tooltip": "Optional external keypoints for garment pose. Used when pose_source is auto/external_pose_keypoints."},
|
|
),
|
|
"person_segmentation_image": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": "Optional external person labelmap. Only affects output when segmentation_free is False and parser_backend is external_fashn_labelmap.",
|
|
},
|
|
),
|
|
"garment_segmentation_image": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": "Optional external garment labelmap. Used when parser_backend is external_fashn_labelmap.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "FashnAI"
|
|
|
|
def process(
|
|
self,
|
|
pipeline,
|
|
person_image,
|
|
garment_image,
|
|
garment_photo_type,
|
|
category,
|
|
skip_cfg_last_n_steps,
|
|
segmentation_free,
|
|
steps,
|
|
cfg,
|
|
seed,
|
|
keep_model_loaded,
|
|
pose_source="auto",
|
|
parser_backend="fashn_human_parser",
|
|
person_pose_keypoints=None,
|
|
garment_pose_keypoints=None,
|
|
person_segmentation_image=None,
|
|
garment_segmentation_image=None,
|
|
):
|
|
device = comfy.model_management.get_torch_device()
|
|
|
|
print(f"FashnVTON: Moving models to {device}...")
|
|
if hasattr(pipeline, "tryon_model"):
|
|
pipeline.tryon_model.to(device)
|
|
if hasattr(pipeline, "hp_model"):
|
|
if hasattr(pipeline.hp_model, "model"):
|
|
pipeline.hp_model.model.to(device)
|
|
pipeline.hp_model.device = device
|
|
|
|
pbar = comfy.utils.ProgressBar(steps)
|
|
|
|
def progress_callback(step, total_steps):
|
|
pbar.update_absolute(step + 1, total_steps)
|
|
|
|
person_pil = _tensor_to_pil(person_image)
|
|
garment_pil = _tensor_to_pil(garment_image)
|
|
|
|
person_segmentation_pil = (
|
|
_tensor_to_pil(person_segmentation_image) if person_segmentation_image is not None else None
|
|
)
|
|
garment_segmentation_pil = (
|
|
_tensor_to_pil(garment_segmentation_image) if garment_segmentation_image is not None else None
|
|
)
|
|
|
|
seed = seed % (2**32)
|
|
|
|
try:
|
|
result = pipeline(
|
|
person_image=person_pil,
|
|
garment_image=garment_pil,
|
|
category=category,
|
|
garment_photo_type=garment_photo_type,
|
|
segmentation_free=segmentation_free,
|
|
skip_cfg_last_n_steps=skip_cfg_last_n_steps,
|
|
num_timesteps=steps,
|
|
guidance_scale=cfg,
|
|
seed=seed,
|
|
pose_source=pose_source,
|
|
person_pose_keypoints=person_pose_keypoints,
|
|
garment_pose_keypoints=garment_pose_keypoints,
|
|
parser_backend=parser_backend,
|
|
person_segmentation_image=person_segmentation_pil,
|
|
garment_segmentation_image=garment_segmentation_pil,
|
|
callback=progress_callback,
|
|
)
|
|
finally:
|
|
# Handle Offloading
|
|
if not keep_model_loaded:
|
|
print("FashnVTON: Unloading models from VRAM...")
|
|
|
|
if hasattr(pipeline, "tryon_model"):
|
|
pipeline.tryon_model.to("cpu")
|
|
if hasattr(pipeline, "hp_model"):
|
|
if hasattr(pipeline.hp_model, "model"):
|
|
pipeline.hp_model.model.to("cpu")
|
|
pipeline.hp_model.device = torch.device("cpu")
|
|
if hasattr(pipeline, "pose_model"):
|
|
del pipeline.pose_model
|
|
if hasattr(pipeline, "pose_provider"):
|
|
pipeline.pose_provider = None
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
# Convert back to ComfyUI format (B, H, W, C)
|
|
output_img = np.array(result.images[0]).astype(np.float32) / 255.0
|
|
output_tensor = torch.from_numpy(output_img).unsqueeze(0)
|
|
|
|
return (output_tensor,)
|
|
|
|
|
|
class FashnPoseKeypointsAdapter:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pose_keypoints": ("POSE_KEYPOINT",),
|
|
"single_person": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FASHN_DWPOSE_KEYPOINTS",)
|
|
FUNCTION = "adapt"
|
|
CATEGORY = "FashnAI/Adapters"
|
|
|
|
def adapt(self, pose_keypoints, single_person):
|
|
converted = convert_pose_keypoints_to_dwpose(pose_keypoints, single_person=single_person)
|
|
if converted is None:
|
|
raise ValueError(
|
|
"Unsupported POSE_KEYPOINT payload for Fashn adapter. "
|
|
"Expected OpenPose-like `people` data or DWPose-style bodies."
|
|
)
|
|
return (converted,)
|
|
|
|
|
|
class FashnMaskToLabelmap:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask_image": ("IMAGE",),
|
|
"category": (["tops", "bottoms", "one-pieces"], {"default": "tops"}),
|
|
"mask_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"label_id_override": (
|
|
"INT",
|
|
{
|
|
"default": -1,
|
|
"min": -1,
|
|
"max": MAX_FASHN_LABEL_ID,
|
|
"step": 1,
|
|
"tooltip": "-1 uses category default label id.",
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "FashnAI/Adapters"
|
|
|
|
def convert(self, mask_image, category, mask_threshold, label_id_override):
|
|
arr = mask_image.detach().cpu().numpy().astype(np.float32)
|
|
if arr.shape[-1] == 1:
|
|
mask = arr[..., 0]
|
|
elif arr.shape[-1] >= 3:
|
|
mask = 0.299 * arr[..., 0] + 0.587 * arr[..., 1] + 0.114 * arr[..., 2]
|
|
else:
|
|
mask = np.mean(arr, axis=-1)
|
|
|
|
mask = np.clip(mask, 0.0, 1.0)
|
|
fg = mask >= float(mask_threshold)
|
|
|
|
label_id = int(label_id_override)
|
|
if label_id < 0:
|
|
label_id = DEFAULT_CATEGORY_LABEL_ID[category]
|
|
|
|
label_map = np.zeros(mask.shape, dtype=np.uint8)
|
|
label_map[fg] = label_id
|
|
|
|
out = np.repeat(label_map[..., None], 3, axis=-1).astype(np.float32) / 255.0
|
|
return (torch.from_numpy(out),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FashnVtonLoader": FashnVtonLoader,
|
|
"FashnVtonInference": FashnVtonInference,
|
|
"FashnPoseKeypointsAdapter": FashnPoseKeypointsAdapter,
|
|
"FashnMaskToLabelmap": FashnMaskToLabelmap,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"FashnVtonLoader": "(Down)load Fashn VTON",
|
|
"FashnVtonInference": "Fashn VTON Inference",
|
|
"FashnPoseKeypointsAdapter": "Fashn Pose Keypoints Adapter",
|
|
"FashnMaskToLabelmap": "Fashn Mask to Labelmap",
|
|
}
|