first commit

2025-03-04 19:11:59 +08:00
parent de84870bc3
commit 727a85df13
6 changed files with 181 additions and 1 deletions
@@ -0,0 +1,6 @@
+_test_*.*
+__pycache__
+.venv
+.idea
+*.pth
+*.ini
@@ -1,2 +1,23 @@
 # ComfyUI_CogView4_Wrapper
-The unofficial implementation of CogView4 project in ComfyUI.
+
+
+The unofficial implementation of [CogView4](https://github.com/THUDM/CogView4) project in ComfyUI.    
+Recommended to run on Nvidia GPU with VRAM of 16GB or more.
+![image](image/cogview4_wrapper_example.jpg)
+
+### Install
+Open the cmd window in the plugin directory of ComfyUI, like ```ComfyUI\custom_nodes```，type  
+```
+git clone https://github.com/chflame163/ComfyUI_CogView4_Wrapper.git
+```
+Install dependency packages：
+```
+pip install -r ComfyUI_CogView4_Wrapper/requirements.txt
+```
+
+The model will be automatically downloaded to the ```ComfyUI/models/CogView``` directory during the first run.
+
+    
+
+## statement
+This nodes follows the MIT license, Some of its functional code comes from other open-source projects. Thanks to the original author. If used for commercial purposes, please refer to the original project license to authorization agreement.
@@ -0,0 +1,48 @@
+import importlib.util
+import os
+import sys
+import json
+
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
+
+python = sys.executable
+
+def get_ext_dir(subpath=None, mkdir=False):
+    dir = os.path.dirname(__file__)
+    if subpath is not None:
+        dir = os.path.join(dir, subpath)
+
+    dir = os.path.abspath(dir)
+
+    if mkdir and not os.path.exists(dir):
+        os.makedirs(dir)
+    return dir
+
+def serialize(obj):
+    if isinstance(obj, (str, int, float, bool, list, dict, type(None))):
+        return obj
+    return str(obj)  # 转为字符串
+
+
+py = get_ext_dir("py")
+files = os.listdir(py)
+all_nodes = {}
+for file in files:
+    if not file.endswith(".py"):
+        continue
+    name = os.path.splitext(file)[0]
+    imported_module = importlib.import_module(".py.{}".format(name), __name__)
+    try:
+        NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **imported_module.NODE_CLASS_MAPPINGS}
+        NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **imported_module.NODE_DISPLAY_NAME_MAPPINGS}
+        serialized_CLASS_MAPPINGS = {k: serialize(v) for k, v in imported_module.NODE_CLASS_MAPPINGS.items()}
+        serialized_DISPLAY_NAME_MAPPINGS = {k: serialize(v) for k, v in imported_module.NODE_DISPLAY_NAME_MAPPINGS.items()}
+        all_nodes[file]={"NODE_CLASS_MAPPINGS": serialized_CLASS_MAPPINGS, "NODE_DISPLAY_NAME_MAPPINGS": serialized_DISPLAY_NAME_MAPPINGS}
+    except:
+        pass
+
+
+# WEB_DIRECTORY = "./js"
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
@@ -0,0 +1,100 @@
+import torch
+import numpy as np
+from PIL import Image
+import os
+import folder_paths
+
+def pil2tensor(image: Image) -> torch.Tensor:
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+def check_and_download_model(model_path, repo_id):
+    model_path = os.path.join(folder_paths.models_dir, "CogView", model_path)
+
+    if not os.path.exists(model_path):
+        print(f"Downloading {repo_id} model to {model_path} ...")
+        from huggingface_hub import snapshot_download
+        snapshot_download(repo_id=repo_id, local_dir=model_path, ignore_patterns=["*.md", "*.txt", ".git", ".gitattributes"])
+    return model_path
+
+class CogView4_Wrapper:
+
+    def __init__(self):
+        self.NODE_NAME = 'CogView4 Wrapper'
+        self.model_name = ""
+        self.dtype = ""
+        self.pipe = None
+
+
+    @classmethod
+    def INPUT_TYPES(self):
+        model_list =['CogView4-6B']
+        dtype_list = ['bf16', 'fp32']
+        default_prompt = "A vibrant cherry red sports car sits proudly under the gleaming sun, its polished exterior smooth and flawless, casting a mirror-like reflection. The car features a low, aerodynamic body, angular headlights that gaze forward like predatory eyes, and a set of black, high-gloss racing rims that contrast starkly with the red. A subtle hint of chrome embellishes the grille and exhaust, while the tinted windows suggest a luxurious and private interior. The scene conveys a sense of speed and elegance, the car appearing as if it's about to burst into a sprint along a coastal road, with the ocean's azure waves crashing in the background.The license plate number of the car is 'CogView4'. The car sprinted along a coastal road, with the same sports car printed on the roadside billboard and large text 'GO AHEAD' written to it. The text was yellow, with thick strokes and heavy shadow lines."
+        return {
+            "required": {
+                "model": (model_list,),
+                "dtype": (dtype_list,),
+                "prompt":("STRING", {"default":default_prompt, "multiline": True}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 1e14}),
+                "guidance_scale": ("FLOAT", {"default": 5, "min": 0.1, "max": 100, "step": 0.1}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
+                "steps": ("INT", {"default": 30, "min": 1, "max": 100, "step": 1}),
+                "width": ("INT", {"default": 1024, "min": 8, "max": 4096, "step": 16}),
+                "height": ("INT", {"default": 1024, "min": 8, "max": 4096, "step": 16}),
+                "cache_model": ("BOOLEAN", {"default": True,}),
+            },
+            "optional": {
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = 'cogview4'
+    CATEGORY = '😺dzNodes/CogView4 Wrapper'
+
+    def cogview4(self, model, prompt, seed, dtype, guidance_scale, batch_size, steps, width, height, cache_model):
+
+        ret_images = []
+
+        from diffusers import CogView4Pipeline
+
+        if self.dtype != dtype or self.model_name != model:
+            model_path = check_and_download_model(model, f"THUDD/{model}")
+            if dtype == 'bf16':
+                self.pipe = CogView4Pipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16).to("cuda")
+                self.model_name = model
+                self.dtype = dtype
+            else:
+                self.pipe = CogView4Pipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cuda")
+                self.model_name = model
+                self.dtype = dtype
+            # Open it for reduce GPU memory usage
+            self.pipe.enable_model_cpu_offload()
+            self.pipe.vae.enable_slicing()
+            self.pipe.vae.enable_tiling()
+
+        image = self.pipe(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=batch_size,
+            num_inference_steps=steps,
+            width=width,
+            height=height,
+            generator=torch.Generator("cuda").manual_seed(seed),
+        ).images
+        for i in image:
+            ret_images.append(pil2tensor(i))
+
+        if not cache_model:
+            self.pipe = None
+            self.model_name = ""
+            self.dtype = ""
+        return (torch.cat(ret_images, dim=0),)
+
+NODE_CLASS_MAPPINGS = {
+    "CogView4": CogView4_Wrapper
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "CogView4": "CogView4 Wrapper"
+}
@@ -0,0 +1,5 @@
+numpy<2.0
+pillow
+torch
+huggingface_hub
+git+https://github.com/huggingface/diffusers.git