diff --git a/README.md b/README.md index 94eec42..7e8110f 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,9 @@ pip install -r custom_nodes/ComfyUI-OOTDiffusion/requirements.txt ## Node 节点 -Load OOTDiffusion: 加载 OOTDiffusion Pipeline +Load OOTDiffusion Local: 从本地目录加载 OOTDiffusion Pipeline + +Load OOTDiffusion from Hub: 从 huggingface 自动下载并加载 OOTDiffusion Pipeline OOTDiffusion Generate: 生成图像 @@ -84,8 +86,10 @@ OOTDiffusion Generate: 生成图像 OOTDiffusion 本体依赖于 `diffusers==0.24.0` 实现，所以假如有其他节点的依赖冲突是没办法解决的（本就不该依赖 diffusers）。 靠 vendor 也能解决，所以也不是大问题。 -不使用 huggingface_hub 是因为 OOTD 提供的仓库并不是一个单纯的 diffusion model structure， -里面还包含了独立的 openpose 和 humanparsing 模型文件。 -目前只有 openai/clip-vit-large-patch14 是使用 huggingface_hub 下载的。 - 在 `Ubuntu 22.02` / `Python 3.10.x` 下可以正常运行。Windows 没有测试过。 + +## 更新日志 Release Note + +2024-02-25: + +移除 git lfs 下载教程，提供 Load OOTDiffusion from Hub 节点 diff --git a/__init__.py b/__init__.py index cb0def4..5707da6 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,7 @@ +import os + import numpy as np +from huggingface_hub import snapshot_download from PIL import Image from torchvision.transforms.functional import to_pil_image, to_tensor @@ -7,7 +10,6 @@ from .inference_ootd import OOTDiffusion from .ootd_utils import get_mask_location from .openpose.run_openpose import OpenPose - _category_get_mask_input = { "upperbody": "upper_body", "lowerbody": "lower_body", @@ -16,11 +18,14 @@ _category_get_mask_input = { class LoadOOTDPipeline: + display_name = "Load OOTDiffusion Local" + @classmethod def INPUT_TYPES(cls): return { "required": { - # "model_type": ("STRING", ["hd", "dc"]), + "type": (["Half body", "Full body"],), + "path": ("STRING", {"default": "models/OOTDiffusion"}), } } @@ -30,11 +35,53 @@ class LoadOOTDPipeline: CATEGORY = "OOTD" - def load(self): - return (OOTDiffusion(hg_root="models/OOTDiffusion"),) + @staticmethod + def load_impl(type, path): + if type == "Half body": + type = "hd" + elif type == "Full body": + type = "dc" + raise RuntimeError("full body is not supported yet") + else: + raise ValueError( + f"unknown input type {type} must be 'Half body' or 'Full body'" + ) + if not os.path.isdir(path): + raise ValueError(f"input path {path} is not a directory") + return OOTDiffusion(path, model_type=type) + + def load(self, type, path): + return (self.load_impl(type, path),) + + +class LoadOOTDPipelineHub(LoadOOTDPipeline): + display_name = "Load OOTDiffusion from Hub🤗" + + repo_id = "levihsu/OOTDiffusion" + repo_revision = "c63b33843e01c8c2c8e591a1d6b88a2feba478b8" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "type": (["Half body", "Full body"],), + } + } + + def load(self, type): # type: ignore + # DiffusionPipeline.from_pretrained doesn't support subfolder + # So we use snapshot_download to get local path first + path = snapshot_download( + self.repo_id, + revision=self.repo_revision, + resume_download=True, + ) + return (LoadOOTDPipeline.load_impl(type, path),) class OOTDGenerate: + display_name = "OOTDiffusion Generate" + @classmethod def INPUT_TYPES(cls): return { @@ -44,6 +91,7 @@ class OOTDGenerate: "model_image": ("IMAGE",), # Openpose from comfyui-controlnet-aux not work # "keypoints": ("POSE_KEYPOINT",), + # TODO: add category when dc model release # "category": ("STRING", ["upperbody", "lowerbody", "dress"]), "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), @@ -116,6 +164,7 @@ class OOTDGenerate: seed=seed, ) + # pil(H,W,3) -> tensor(H,W,3) output_image = to_tensor(images[0]) output_image = output_image.permute((1, 2, 0)) masked_vton_img = masked_vton_img.convert("RGB") @@ -125,12 +174,14 @@ class OOTDGenerate: return ([output_image], [masked_vton_img]) -NODE_CLASS_MAPPINGS = { - "LoadOOTDPipeline": LoadOOTDPipeline, - "OOTDGenerate": OOTDGenerate, -} +_export_classes = [ + LoadOOTDPipeline, + LoadOOTDPipelineHub, + OOTDGenerate, +] + +NODE_CLASS_MAPPINGS = {c.__name__: c for c in _export_classes} NODE_DISPLAY_NAME_MAPPINGS = { - "LoadOOTDPipeline": "Load OOTDiffusion", - "OOTDGenerate": "OOTDiffusion Generate", + c.__name__: getattr(c, "display_name", c.__name__) for c in _export_classes } diff --git a/inference_ootd.py b/inference_ootd.py index d1e682a..484a8b1 100644 --- a/inference_ootd.py +++ b/inference_ootd.py @@ -1,14 +1,11 @@ +import os import random import sys import time from pathlib import Path -import cv2 -import numpy as np import torch -import torch.nn.functional as F from diffusers import AutoencoderKL, UniPCMultistepScheduler -from PIL import Image from transformers import ( AutoProcessor, CLIPTextModel, @@ -16,7 +13,6 @@ from transformers import ( CLIPVisionModelWithProjection, ) -# sys.path.append(str(Path(__file__).parent)) from . import pipelines_ootd #! Necessary for OotdPipeline.from_pretrained @@ -29,7 +25,7 @@ from .pipelines_ootd.unet_vton_2d_condition import UNetVton2DConditionModel class OOTDiffusion: - def __init__(self, hg_root: str, model_type: str = "hd"): + def __init__(self, root: str, model_type: str = "hd"): self.device = "cuda" if torch.cuda.is_available() else "cpu" if model_type not in ("hd", "dc"): @@ -37,20 +33,13 @@ class OOTDiffusion: self.model_type = model_type - # hg_root = str(Path.cwd().resolve() / hg_root) VIT_PATH = f"openai/clip-vit-large-patch14" - VAE_PATH = f"{hg_root}/checkpoints/ootd" + MODEL_PATH = Path(f"{root}/checkpoints/ootd") if model_type == "hd": - UNET_PATH = f"{hg_root}/checkpoints/ootd/ootd_hd/checkpoint-36000" + UNET_PATH = MODEL_PATH / "ootd_hd" / "checkpoint-36000" else: - UNET_PATH = f"{hg_root}/checkpoints/ootd/ootd_dc/checkpoint-36000" - MODEL_PATH = f"{hg_root}/checkpoints/ootd" + UNET_PATH = MODEL_PATH / "ootd_dc" / "checkpoint-36000" - vae = AutoencoderKL.from_pretrained( - VAE_PATH, - subfolder="vae", - torch_dtype=torch.float16, - ) unet_garm = UNetGarm2DConditionModel.from_pretrained( UNET_PATH, subfolder="unet_garm", @@ -67,7 +56,10 @@ class OOTDiffusion: MODEL_PATH, unet_garm=unet_garm, unet_vton=unet_vton, - vae=vae, + vae=AutoencoderKL.from_pretrained( + f"{MODEL_PATH}/vae", + torch_dtype=torch.float16, + ), torch_dtype=torch.float16, variant="fp16", use_safetensors=True, @@ -82,14 +74,8 @@ class OOTDiffusion: self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(VIT_PATH).to( self.device ) - self.tokenizer = CLIPTokenizer.from_pretrained( - MODEL_PATH, - subfolder="tokenizer", - ) - self.text_encoder = CLIPTextModel.from_pretrained( - MODEL_PATH, - subfolder="text_encoder", - ).to(self.device) + self.tokenizer = self.pipe.tokenizer + self.text_encoder = self.pipe.text_encoder def tokenize_captions(self, captions, max_length): inputs = self.tokenizer(