add node Load OOTDiffusion from Hub

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 节点

__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
 }

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(