diff --git a/README.md b/README.md index 090d2a7..9983eb6 100644 --- a/README.md +++ b/README.md @@ -102,6 +102,10 @@ OOTDiffusion 本体依赖于 `diffusers==0.24.0` 实现,所以假如有其他 ## 更新日志 Release Note +2024-03-04: + +添加 Full body 模型 + 2024-03-01: 添加 Windows 安装教程 diff --git a/__init__.py b/__init__.py index ac7dbf2..a1f8e1d 100644 --- a/__init__.py +++ b/__init__.py @@ -16,6 +16,12 @@ _category_get_mask_input = { "dress": "dresses", } +_category_readable = { + "Upper body": "upperbody", + "Lower body": "lowerbody", + "Dress": "dress", +} + class LoadOOTDPipeline: display_name = "Load OOTDiffusion Local" @@ -41,7 +47,6 @@ class LoadOOTDPipeline: 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'" @@ -58,7 +63,7 @@ class LoadOOTDPipelineHub(LoadOOTDPipeline): display_name = "Load OOTDiffusion from Hub🤗" repo_id = "levihsu/OOTDiffusion" - repo_revision = "c63b33843e01c8c2c8e591a1d6b88a2feba478b8" + repo_revision = "6150dd90d7302f21348b8d8461766a400c00716a" @classmethod def INPUT_TYPES(cls): @@ -96,8 +101,6 @@ 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}), "cfg": ( @@ -110,6 +113,7 @@ class OOTDGenerate: "round": 0.01, }, ), + "category": (list(_category_readable.keys()),), } } @@ -119,8 +123,9 @@ class OOTDGenerate: CATEGORY = "OOTD" - def generate(self, pipe, cloth_image, model_image, seed, steps, cfg): - category = "upperbody" + def generate( + self, pipe: OOTDiffusion, cloth_image, model_image, category, seed, steps, cfg + ): # if model_image.shape != (1, 1024, 768, 3) or ( # cloth_image.shape != (1, 1024, 768, 3) # ): @@ -128,6 +133,11 @@ class OOTDGenerate: # f"Input image must be size (1, 1024, 768, 3). " # f"Got model_image {model_image.shape} cloth_image {cloth_image.shape}" # ) + category = _category_readable[category] + if pipe.model_type == "hd" and category != "upperbody": + raise ValueError( + "Half body (hd) model type can only be used with upperbody category" + ) # (1,H,W,3) -> (3,H,W) model_image = model_image.squeeze(0)