add onediff_comfy_nodes/benchmarks

.github/workflows/examples.yml

@@ -282,7 +282,7 @@ jobs:
       - name: Start ComfyUI Web Service
         if: matrix.test-suite == 'comfy'
         run: |
-          docker exec -w /app/ComfyUI -d ${{ env.CONTAINER_NAME }} sh -c "python3 /app/ComfyUI/main.py --port 8188 --extra-model-paths-config /src/onediff/tests/comfyui/extra_model_paths.yaml > /app/ComfyUI/onediff_comfyui.log 2>&1"
+          docker exec -w /app/ComfyUI -d ${{ env.CONTAINER_NAME }} sh -c "python3 /app/ComfyUI/main.py --gpu-only --disable-cuda-malloc --port 8188 --extra-model-paths-config /src/onediff/tests/comfyui/extra_model_paths.yaml > /app/ComfyUI/onediff_comfyui.log 2>&1"
           sleep 30
       # print to check if comfy is launched successfully
       - run: docker exec ${{ env.CONTAINER_NAME }} ps aux
@@ -298,13 +298,15 @@ jobs:
               false
             }
           }
-
-          run_comfy_test "workflows/sdxl-unet-speedup-graph-saver.json" 200
-          run_comfy_test "workflows/sdxl-control-lora-speedup.json" 200
-          run_comfy_test "/share_nfs/hf_models/comfyui_resources/workflows/ipadapter_advanced.json" 200
-          run_comfy_test "/share_nfs/hf_models/comfyui_resources/workflows/deep-cache.json" 600
-          run_comfy_test "/share_nfs/hf_models/comfyui_resources/workflows/deep-cache-with-lora.json" 800
+          # run_comfy_test "/share_nfs/hf_models/comfyui_resources/workflows/deep-cache.json" 600
+          # run_comfy_test "/share_nfs/hf_models/comfyui_resources/workflows/deep-cache-with-lora.json" 800
           # run_comfy_test "workflows/text-to-video-speedup.json" 5000
+          docker exec -w /src/onediff/onediff_comfy_nodes/benchmarks ${{ env.CONTAINER_NAME }} bash scripts/install_env.sh /app/ComfyUI
+          docker exec -w /src/onediff/onediff_comfy_nodes/benchmarks ${{ env.CONTAINER_NAME }} bash scripts/run_all_tests.sh || {
+              echo "Test fails! print the ComfyUI logs..."
+              docker exec onediff-test cat /app/ComfyUI/onediff_comfyui.log
+              false
+            }
 
       - name: Show ComfyUI Log
         if: matrix.test-suite == 'comfy'

.gitignore

@@ -180,3 +180,4 @@ unet_graphs
 
 # onediff_sd_webui_extensions
 onediff_sd_webui_extensions/compiled_caches/
+onediff_comfy_nodes/benchmarks/results/

onediff_comfy_nodes/_config.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import torch
 import folder_paths
 
 __all__ = [
@@ -8,6 +9,7 @@
     "is_disable_oneflow_backend",
 ]
 
+
 # https://github.com/comfyanonymous/ComfyUI/blob/master/folder_paths.py#L9
 os.environ["COMFYUI_ROOT"] = folder_paths.base_path
 _default_backend = os.environ.get("ONEDIFF_COMFY_NODES_DEFAULT_BACKEND", "oneflow")

onediff_comfy_nodes/_nodes.py

@@ -13,10 +13,12 @@
     from .modules.oneflow import BasicOneFlowBoosterExecutor
 
     BasicBoosterExecutor = BasicOneFlowBoosterExecutor
+    print("\033[1;31mUsing OneFlow backend\033[0m (Default)")
 elif is_nexfort_available():
     from .modules.nexfort.booster_basic import BasicNexFortBoosterExecutor
 
     BasicBoosterExecutor = BasicNexFortBoosterExecutor
+    print("\033[1;32mUsing Nexfort backend\033[0m (Default)")
 else:
     raise RuntimeError(
         "Neither OneFlow nor Nexfort is available. Please ensure at least one of them is installed."
@@ -44,6 +46,7 @@ def speedup(
         model,
         inplace: bool = False,
         custom_booster: Optional[BoosterScheduler] = None,
+        booster_settings: Optional[BoosterSettings] = None,
         *args,
         **kwargs
     ) -> Tuple:
@@ -60,15 +63,17 @@ def speedup(
         Returns:
             Tuple: Tuple containing the optimized model.
         """
-        if not hasattr(self, "booster_settings"):
+        if booster_settings is None and not hasattr(self, "booster_settings"):
             self.booster_settings = BoosterSettings(tmp_cache_key=str(uuid.uuid4()))
 
         if custom_booster:
             booster = custom_booster
             booster.inplace = inplace
         else:
             booster = BoosterScheduler(BasicBoosterExecutor(), inplace=inplace)
-        booster.settings = self.booster_settings
+        booster.settings = (
+            self.booster_settings if booster_settings is None else booster_settings
+        )
         return (booster(model, *args, **kwargs),)
 
 
@@ -93,6 +98,9 @@ def INPUT_TYPES(s):
 
     RETURN_TYPES = ("VAE",)
 
+    def speedup(self, vae, inplace=False, custom_booster: BoosterScheduler = None):
+        return super().speedup(vae, inplace, custom_booster)
+
 
 class ControlnetSpeedup:
     @classmethod
@@ -193,7 +201,7 @@ def onediff_load_controlnet(self, control_net_name, custom_booster=None):
         return (controlnet,)
 
 
-class OneDiffCheckpointLoaderSimple(CheckpointLoaderSimple):
+class OneDiffCheckpointLoaderSimple(CheckpointLoaderSimple, SpeedupMixin):
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -207,40 +215,35 @@ def INPUT_TYPES(s):
     CATEGORY = "OneDiff/Loaders"
     FUNCTION = "onediff_load_checkpoint"
 
-    @staticmethod
-    def _load_checkpoint(
-        ckpt_name, vae_speedup="disable", custom_booster: BoosterScheduler = None
-    ):
-        """Loads a checkpoint, applying speedup techniques."""
-
-        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
-        out = comfy.sd.load_checkpoint_guess_config(
-            ckpt_path,
-            output_vae=True,
-            output_clip=True,
-            embedding_directory=folder_paths.get_folder_paths("embeddings"),
-        )
-
-        # Unpack outputs
-        modelpatcher, clip, vae = out[:3]
+    def __init__(self) -> None:
+        super().__init__()
+        self.unet_booster_settings = BoosterSettings(tmp_cache_key=str(uuid.uuid4()))
+        self.vae_booster_settings = BoosterSettings(tmp_cache_key=str(uuid.uuid4()))
 
-        # Apply custom booster if provided, otherwise use a basic one
-        custom_booster = custom_booster or BoosterScheduler(BasicBoosterExecutor())
-        modelpatcher = custom_booster(modelpatcher, ckpt_name=ckpt_name)
+    @torch.inference_mode()
+    def onediff_load_checkpoint(
+        self, ckpt_name, vae_speedup="disable", custom_booster: BoosterScheduler = None,
+    ):
+        modelpatcher, clip, vae = self.load_checkpoint(ckpt_name)
+        modelpatcher = self.speedup(
+            modelpatcher,
+            inplace=True,
+            custom_booster=custom_booster,
+            booster_settings=self.unet_booster_settings,
+        )[0]
 
-        # Apply VAE speedup if enabled
         if vae_speedup == "enable":
-            vae = BoosterScheduler(BasicBoosterExecutor())(vae, ckpt_name=ckpt_name)
+            vae = self.speedup(
+                vae,
+                inplace=True,
+                custom_booster=custom_booster,
+                booster_settings=self.vae_booster_settings,
+            )[0]
 
         # Set weight inplace update
         modelpatcher.weight_inplace_update = True
-
-        return modelpatcher, clip, vae
-
-    @torch.inference_mode()
-    def onediff_load_checkpoint(
-        self, ckpt_name, vae_speedup="disable", custom_booster: BoosterScheduler = None,
-    ):
-        out = self._load_checkpoint(ckpt_name, vae_speedup, custom_booster)
-        # Return the loaded checkpoint (modelpatcher, clip, vae)
-        return out
+        return (
+            modelpatcher,
+            clip,
+            vae,
+        )

onediff_comfy_nodes/benchmarks/README.md

@@ -0,0 +1,59 @@
+## Environment setup
+### Set up ComfyUI
+https://github.com/comfyanonymous/ComfyUI
+
+### Set up onediff_comfy_nodes
+https://github.com/siliconflow/onediff?tab=readme-ov-file#installation
+
+```shell
+# python 3.10
+git clone https://github.com/siliconflow/onediff.git
+cd onediff && pip install -e .
+ln -s $(pwd)/onediff_comfy_nodes path/to/ComfyUI/custom_nodes/
+# or
+# cp -r onediff_comfy_nodes path/to/ComfyUI/custom_nodes/
+```
+### Set up nexfort backend
+https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
+
+## Getting Started
+### Run ComfyUI
+Note ⚠️: Replace 'path/to/' with the actual path to the directories and files on your system.
+```shell
+export COMFYUI_ROOT=path/to/ComfyUI
+
+cd path/to/onediff/onediff_comfy_nodes/benchmarks
+
+bash scripts/install_env.sh $COMFYUI_ROOT
+
+cd  $COMFYUI_ROOT
+
+python main.py --gpu-only --port 8188 --extra-model-paths-config path/to/onediff/tests/comfyui/extra_model_paths.yaml 
+```
+
+## Usage Example
+
+```shell
+cd path/to/onediff/onediff_comfy_nodes/benchmarks
+
+bash scripts/run_text_to_image.sh
+```
+
+The output results will be saved in the results/ directory.
+
+## How to add a workflow for testing
+To add a workflow for testing, you can refer to the `run_text_to_image.sh` script and the `input_registration.py` file in the `src` directory. Here's an example of how to register a workflow generator:
+
+```python
+# file: onediff/onediff_comfy_nodes/benchmarks/src/input_registration.py
+
+@register_generator(f"{WORKFLOW_DIR}/example_workflow_api.json")
+def _(workflow_path, *args, **kwargs):
+    with open(workflow_path, "r") as fp:
+        workflow = json.load(fp)
+    graph = ComfyGraph(graph=workflow, sampler_nodes=["3"])
+    for height in [1024, 768, 512]:
+        for width in [1024, 768, 512]:
+            graph.set_image_size(height=height, width=width)
+            yield InputParams(graph=graph)
+```

onediff_comfy_nodes/benchmarks/resources/prompts.txt

@@ -0,0 +1,81 @@
+universe,stars,moon
+Morning fog over sleepy village.
+A group of giraffes eating leaves，sweetveld
+in style of Kawanabe Kyosai, beautiful details
+a dog,hold hot dog,outdoors,grass
+A pickup truck going up a mountain switchback
+texture of marble, top down close-up, video game
+robot droids, in the desert , colorful, dutch angle
+Man snowboarding on the mars, ultra high resolution 8k
+ink cartoon frazzled cute cat, nervous cat, white background
+texture of wood bark, top down close-up, video game
+1boy,underwater,green eyes,white skirt,looking at viewer
+A vibrant tropical rainforest, cold color palette, muted colors, detailed
+Fluffy Samoye, smile, holding a toy ball in his mouth
+in style of Alphonso Mucha , character concept design, half body
+Cyberpunk style,urban,1 robot,an electronic screen with “Khazix”
+in style of E.H. Shepard , character, ink art, side view
+breathtaking night street of Tokyo, neon lights. award-winning, professional, highly detailed
+concept art {prompt}. digital artwork, illustrative, painterly, matte painting, highly detailed
+16-bit pixel art, a cozy cafe side view, a beautiful day
+Byzantine Mosaic Art of a single purple flower in a withe vase.
+A cute cat with a sign saying "Go Big or Go Home".
+a cute cat with a sign saying "Go Big or Go Home"
+A landscape photo of Iceland, with aurora, snow, ice and erupting lava
+Fauvist Depiction of a Sunlit Village with Simplified Forms and Intense Color Contrasts.
+A majestic lion stands proudly on a rock, overlooking the vast African savannah
+anime artwork an empty classroom. anime style, key visual, vibrant, studio anime, highly detailed
+claymation style captain jack sparrow on tropical island. sculpture, clay art, centered composition, play-doh
+claymation style captain jack sparrow on tropical island. sculpture, clay art, centered composition, play-doh
+isometric pixel-art of wizard working on spells. low-res, blocky, pixel art style, 16-bit graphics
+a latina woman with a pearl earring,best quality,masterpiece,ultra detailed,UHD 4K,photographic
+isometric style farmhouse from RPG game, unreal engine, vibrant, beautiful, crisp, detailed, ultra detailed, intricate
+professional 3d model bulky purple mecha with missiles. octane render, highly detailed, volumetric, dramatic lighting
+The 90s, a beautiful woman with a radiant smile and long hair, dressed in summer attire
+A dolphin leaps through the waves, set against a backdrop of bright blues and teal hues
+breathtaking selfie photograph of astronaut floating in space, earth in the background. award-winning, professional, highly detailed
+origami style Winterfell. paper art, pleated paper, folded, origami art, pleats, cut and fold, centered composition
+line art of a kitchen in perspective. professional, sleek, modern, minimalist, graphic, line art, vector graphics
+High resolution HDR photograph of a broken heart, crying by artist "Volko Merschky", by artist"Simone Pfaff".
+A cute rusty robot made if garbage with a panel "I'M SD3-BASED" above it, an evil pupeeter.
+a female character design, short hair wearing headphones, background in the city,Pixel style, transparent pvc jacket
+cinematic photo of a woman sitting at a cafe. 35mm photograph, film, bokeh, professional, 4k, highly detailed
+a man wearing a corduroy outfit at a fashion photoshoot, best quality,masterpiece,ultra detailed,UHD 4K,photographic
+anime artwork a girl looking at the sea, dramatic, anime style, key visual, vibrant, studio anime, highly detailed
+concept art of dragon flying over town, clouds. digital artwork, illustrative, painterly, matte painting, highly detailed, cinematic composition
+cinematic photo of a construction worker looking down at city. 35mm photograph, film, bokeh, professional, 4k, highly detailed
+Anime girl, a beautiful warrior with flowing silver hair stands in the midst of a battle - ravaged landscape.
+vaporwave synthwave style Los Angeles street. cyberpunk, neon, vibes, stunningly beautiful, crisp, detailed, sleek, ultramodern, high contrast, cinematic composition
+low-poly style tropical island with clouds, ambient occlusion . low-poly game art, polygon mesh, jagged, blocky, wireframe edges, centered composition
+a cat,a destroyed badly damaged space ship,beautiful beach,broken windows, grass and flowers grow around,sunny,ocean
+top down shot of a miniature cottage with garden, delicate details, vivid scene, lovely style, cute, sweet, soft atmosphere, Microlandscape
+angel prays in the red shore flower, in the style of David Hettinger, white and red, Dmitry Kustanovich, vibrant statues
+Vincent Willem van Gogh，Decorate with bright and exaggerated col，High brightness， high purity， and high brightness colors，Modern and fashionable color decoration
+ethereal fantasy concept art of sorceress casting spells. magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy
+ethereal fantasy concept art of thunder god with hammer. magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy
+An ornate, Victorian-era key lying on a weathered, wooden surface, with intricate, steampunk-inspired gears and mechanisms visible within its transparent, glass shaft.
+A quirky, steampunk robot with brass gears and a top hat, serving tea in a Victorian parlor, captured in a whimsical photorealistic style.
+A beautiful painting of flowing colors and styles forming the words “SD3 is coming!”, the background is speckled with drops and splashes of paint.
+In the shadow of the last sun, a fisherman had fallen asleep, and he had a furrow along his face, like a sort of smile.
+cinematic film still, stormtrooper taking aim. shallow depth of field, vignette, highly detailed, high budget Hollywood movie, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy
+Inside the cockpit of an airplane during sunset, a futuristic city is visible in the distance. A faint crescent moon can be seen in the sky.
+analog film photo of old woman on the streets of london . faded film, desaturated, 35mm photo, grainy, vignette, vintage, Kodachrome, Lomography, stained, highly detailed, found footage
+Photo of a bear wearing a suit and tophat in a river in the middle of a forest holding a sign that says"I can't bear it".
+An anthopomorphic pink donut with a mustache and cowboy hat standing by a log cabin in a forest with an old 1970s orange truck in the driveway.
+An elegant and feminine packaging design for a luxury, scented candle collection, with a soft, pastel color palette, a frosted glass jar, and a delicate, floral-patterned box.
+Awesome artwork of a wizard on the top of a mountain, he's creating the big text "Stable Diffusion 3 API" with magic, magic text, at dawn, sunrise.
+vaporwave style of porsche driving through tokyo. cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional
+A brooding male character with tousled black hair and a cape stands with a determined expression, backlit by fiery, swirling brushstrokes of orange and yellow, creating a dramatic and intense mood.
+A dark-armored warrior with ornate golden details, cloaked in a flowing black cape, wielding a radiant, fiery sword, standing amidst an ominous cloudy backdrop with dramatic lighting, exuding a menacing, powerful presence.
+Side angle of a Chinese woman wearing a black coat, 25 years old, long straight black hair, street shot, in winter, dynamic blur, black and white, retro, side shot, side, panorama, HD, 16K
+((anime style)),1girl, indoors, sitting on the sofa, living room, pink hair, blue eyes, from back, from above, face towards viewer, playing video games, holding controller, white shirt, short, parted lips, anime production
+A black and sleek robot is looking at a glass blue orb full of electrical current hovering in front of him. The robot has red glowing eyes. The background has a soft glowing red light.
+Cartoon hand, little girl, long colored hair, holographic coat, white shorts, fair skin, blue eyes, white sneakers, full-body photo, full body, panorama, best quality, best picture quality, black highly reflective background cloth, movie-level lighting effect
+Masterpiece, best quality, girl, having a tattoo that says "Welcome to SiliconFlow". collarbone, wavy hair, looking at viewer, blurry foreground, upper body, necklace, contemporary, plain pants, intricate, print, pattern, ponytail, red hair, dappled sunlight, smile, happy.
+A classical painting of a gothic black knight standing in the middle of a field of white flowers. Tall trees are in the distance. The lighting is somber and dramatic, during the golden hour of the day.
+A portrait painting of a handsome man with dark hair and dreamy green eyes. His hair is dark brown, curly, and short. The background is dark and moody with hints of red. The lighting is somber and dramatic.
+line art, line style, 1girl, solo, japanese clothes, hand fan, kimono, black hair, paper fan, outdoors, holding, holding fan, upper body, floral print, short hair, black eyes, brown eyes, black kimono, yukata, blue kimono, uchiwa, closed mouth, sash, expressionless, medium hair
+A boy with a sword in his hand takes a fighting stance on the high street with a sword slashed at the camera, close-up, pov view, first person, blue whirlwind flame, glow, surrealism, ultra-futurism, cyberpunk, 3D art, rich detail, best quality, centered
+Anime girl, masterpiece, best quality, hatsune miku, white gown, angel, angel wings, golden halo, dark background, upper body, closed mouth, looking at viewer, arms behind back, blue theme, night, highres, 4k, 8k, intricate detail, cinematic lighting, amazing quality, amazing shading, soft lighting, detailed Illustration, anime style, wallpaper.
+Super close-up shot of a weathered male skull almost buried by sand,side view,a fresh plant with two green leaves growing from the skull,detailed texture,shot from the botton,epic,super photorealism,cinematic,scenery,sunset,wasteland,desert,dune wave,super-detailed,highly realistic,8k,artistic,contrast lighting,vibrant color,hdr,erode
+a (side view close-up half-body:1.85) fashion photoshoot photo of darth vader wearing a pink and white diamond studded outfit, his chest has a (very big CRT screen showing a pacman game:1.7), his helmet is made of a hello kitty themed white plastic, his helmet has sticker decals on it