From 88b8744096368aa9ff30c8e0a2361f2c2723069b Mon Sep 17 00:00:00 2001 From: Alex Dodge Date: Mon, 27 Feb 2023 18:56:47 -0800 Subject: [PATCH] ... --- .../src/nodes/impl/stable_diffusion/types.py | 26 ++++++++++++++++--- .../nodes/nodes/stable_diffusion/k_sampler.py | 16 ++++++------ .../stable_diffusion/k_sampler_advanced.py | 16 ++++++------ .../nodes/stable_diffusion/latent_upscale.py | 16 +++++------- .../nodes/stable_diffusion/load_model.py | 17 +++++++----- .../nodes/stable_diffusion/vae_decode.py | 5 ++-- .../nodes/stable_diffusion/vae_encode.py | 4 +-- .../stable_diffusion/vae_masked_encode.py | 6 ++--- 8 files changed, 61 insertions(+), 45 deletions(-) diff --git a/backend/src/nodes/impl/stable_diffusion/types.py b/backend/src/nodes/impl/stable_diffusion/types.py index 03103b2e51..2865279bea 100644 --- a/backend/src/nodes/impl/stable_diffusion/types.py +++ b/backend/src/nodes/impl/stable_diffusion/types.py @@ -1,9 +1,27 @@ -from comfy import CLIPModel, Conditioning, LatentImage, StableDiffusionModel, VAEModel +from comfy.clip import CLIPModel +from comfy.conditioning import Conditioning +from comfy.latent_image import CropMethod, LatentImage, UpscaleMethod +from comfy.stable_diffusion import ( + BuiltInCheckpointConfigName, + CheckpointConfig, + Sampler, + Scheduler, + StableDiffusionModel, + load_checkpoint, +) +from comfy.vae import VAEModel __all__ = [ - "StableDiffusionModel", - "VAEModel", "CLIPModel", - "LatentImage", "Conditioning", + "CropMethod", + "LatentImage", + "UpscaleMethod", + "BuiltInCheckpointConfigName", + "CheckpointConfig", + "Sampler", + "Scheduler", + "StableDiffusionModel", + "load_checkpoint", + "VAEModel", ] diff --git a/backend/src/nodes/nodes/stable_diffusion/k_sampler.py b/backend/src/nodes/nodes/stable_diffusion/k_sampler.py index a751c98902..2086ab68e4 100644 --- a/backend/src/nodes/nodes/stable_diffusion/k_sampler.py +++ b/backend/src/nodes/nodes/stable_diffusion/k_sampler.py @@ -1,10 +1,10 @@ from __future__ import annotations -import comfy - from ...impl.stable_diffusion.types import ( Conditioning, LatentImage, + Sampler, + Scheduler, StableDiffusionModel, ) from ...node_base import NodeBase, group @@ -43,12 +43,12 @@ def __init__(self): ), SliderInput("Steps", minimum=1, default=20, maximum=150), EnumInput( - comfy.Sampler, - default_value=comfy.Sampler.SAMPLE_EULER, + Sampler, + default_value=Sampler.SAMPLE_EULER, ), EnumInput( - comfy.Scheduler, - default_value=comfy.Scheduler.NORMAL, + Scheduler, + default_value=Scheduler.NORMAL, ), SliderInput( "CFG Scale", @@ -77,8 +77,8 @@ def run( denoising_strength: float, seed: int, steps: int, - sampler: comfy.Sampler, - scheduler: comfy.Scheduler, + sampler: Sampler, + scheduler: Scheduler, cfg_scale: float, ) -> LatentImage: img = model.sample( diff --git a/backend/src/nodes/nodes/stable_diffusion/k_sampler_advanced.py b/backend/src/nodes/nodes/stable_diffusion/k_sampler_advanced.py index 4601d1cf06..f8028603cd 100644 --- a/backend/src/nodes/nodes/stable_diffusion/k_sampler_advanced.py +++ b/backend/src/nodes/nodes/stable_diffusion/k_sampler_advanced.py @@ -1,10 +1,10 @@ from __future__ import annotations -import comfy - from ...impl.stable_diffusion.types import ( Conditioning, LatentImage, + Sampler, + Scheduler, StableDiffusionModel, ) from ...node_base import NodeBase, group @@ -43,12 +43,12 @@ def __init__(self): ), SliderInput("Steps", minimum=1, default=20, maximum=150), EnumInput( - comfy.Sampler, - default_value=comfy.Sampler.SAMPLE_EULER, + Sampler, + default_value=Sampler.SAMPLE_EULER, ), EnumInput( - comfy.Scheduler, - default_value=comfy.Scheduler.NORMAL, + Scheduler, + default_value=Scheduler.NORMAL, ), SliderInput( "CFG Scale", @@ -91,8 +91,8 @@ def run( denoising_strength: float, seed: int, steps: int, - sampler: comfy.Sampler, - scheduler: comfy.Scheduler, + sampler: Sampler, + scheduler: Scheduler, cfg_scale: float, start_at: int, end_at: int, diff --git a/backend/src/nodes/nodes/stable_diffusion/latent_upscale.py b/backend/src/nodes/nodes/stable_diffusion/latent_upscale.py index d2743025f7..d479d0d271 100644 --- a/backend/src/nodes/nodes/stable_diffusion/latent_upscale.py +++ b/backend/src/nodes/nodes/stable_diffusion/latent_upscale.py @@ -1,8 +1,6 @@ from __future__ import annotations -import comfy.latent_image - -from ...impl.stable_diffusion.types import LatentImage +from ...impl.stable_diffusion.types import CropMethod, LatentImage, UpscaleMethod from ...node_base import NodeBase from ...node_factory import NodeFactory from ...properties.inputs import EnumInput, SliderInput @@ -19,12 +17,12 @@ def __init__(self): self.inputs = [ LatentImageInput(), EnumInput( - comfy.UpscaleMethod, - default_value=comfy.UpscaleMethod.BILINEAR, + UpscaleMethod, + default_value=UpscaleMethod.BILINEAR, ), EnumInput( - comfy.CropMethod, - default_value=comfy.CropMethod.DISABLED, + CropMethod, + default_value=CropMethod.DISABLED, ), SliderInput( "width", @@ -57,8 +55,8 @@ def __init__(self): def run( self, latent_image: LatentImage, - upscale_method: comfy.latent_image.UpscaleMethod, - crop_method: comfy.latent_image.CropMethod, + upscale_method: UpscaleMethod, + crop_method: CropMethod, width: int, height: int, ) -> LatentImage: diff --git a/backend/src/nodes/nodes/stable_diffusion/load_model.py b/backend/src/nodes/nodes/stable_diffusion/load_model.py index ca67114397..6336eef00c 100644 --- a/backend/src/nodes/nodes/stable_diffusion/load_model.py +++ b/backend/src/nodes/nodes/stable_diffusion/load_model.py @@ -3,9 +3,14 @@ import os from typing import Tuple -import comfy - -from ...impl.stable_diffusion.types import CLIPModel, StableDiffusionModel, VAEModel +from ...impl.stable_diffusion.types import ( + BuiltInCheckpointConfigName, + CheckpointConfig, + CLIPModel, + StableDiffusionModel, + VAEModel, + load_checkpoint, +) from ...node_base import NodeBase from ...node_factory import NodeFactory from ...properties.inputs import CkptFileInput @@ -45,11 +50,9 @@ def run( assert os.path.isfile(path), f"Path {path} is not a file" # TODO load V2 models, maybe auto-detect - config = comfy.CheckpointConfig.from_built_in( - comfy.BuiltInCheckpointConfigName.V1 - ) + config = CheckpointConfig.from_built_in(BuiltInCheckpointConfigName.V1) - sd, clip, vae = comfy.load_checkpoint( + sd, clip, vae = load_checkpoint( config=config, checkpoint_filepath=path, embedding_directory=None ) diff --git a/backend/src/nodes/nodes/stable_diffusion/vae_decode.py b/backend/src/nodes/nodes/stable_diffusion/vae_decode.py index 65be884bdd..33b47ec5b9 100644 --- a/backend/src/nodes/nodes/stable_diffusion/vae_decode.py +++ b/backend/src/nodes/nodes/stable_diffusion/vae_decode.py @@ -1,10 +1,9 @@ from __future__ import annotations -import comfy import cv2 import numpy as np -from ...impl.stable_diffusion.types import LatentImage +from ...impl.stable_diffusion.types import LatentImage, VAEModel from ...node_base import NodeBase from ...node_factory import NodeFactory from ...properties.inputs.stable_diffusion_inputs import LatentImageInput, VAEModelInput @@ -30,7 +29,7 @@ def __init__(self): self.icon = "PyTorch" self.sub = "Input & Output" - def run(self, vae: comfy.VAEModel, latent_image: LatentImage) -> np.ndarray: + def run(self, vae: VAEModel, latent_image: LatentImage) -> np.ndarray: img = vae.decode(latent_image) arr = np.array(img) arr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR) diff --git a/backend/src/nodes/nodes/stable_diffusion/vae_encode.py b/backend/src/nodes/nodes/stable_diffusion/vae_encode.py index 07d8f34809..61b091baa5 100644 --- a/backend/src/nodes/nodes/stable_diffusion/vae_encode.py +++ b/backend/src/nodes/nodes/stable_diffusion/vae_encode.py @@ -1,9 +1,9 @@ from __future__ import annotations -import comfy import numpy as np from PIL import Image +from ...impl.stable_diffusion.types import LatentImage, VAEModel from ...node_base import NodeBase from ...node_factory import NodeFactory from ...properties.inputs import ImageInput @@ -35,7 +35,7 @@ def __init__(self): self.icon = "PyTorch" self.sub = "Input & Output" - def run(self, vae: comfy.VAEModel, image: np.ndarray) -> np.ndarray: + def run(self, vae: VAEModel, image: np.ndarray) -> LatentImage: img = _array_to_image(image) latent = vae.encode(img) return latent diff --git a/backend/src/nodes/nodes/stable_diffusion/vae_masked_encode.py b/backend/src/nodes/nodes/stable_diffusion/vae_masked_encode.py index 99ffafc45e..1b10fd825e 100644 --- a/backend/src/nodes/nodes/stable_diffusion/vae_masked_encode.py +++ b/backend/src/nodes/nodes/stable_diffusion/vae_masked_encode.py @@ -1,9 +1,9 @@ from __future__ import annotations -import comfy import numpy as np from PIL import Image +from ...impl.stable_diffusion.types import LatentImage, VAEModel from ...node_base import NodeBase from ...node_factory import NodeFactory from ...properties.inputs import ImageInput @@ -31,9 +31,7 @@ def __init__(self): self.icon = "PyTorch" self.sub = "Input & Output" - def run( - self, vae: comfy.VAEModel, image: np.ndarray, mask: np.ndarray - ) -> np.ndarray: + def run(self, vae: VAEModel, image: np.ndarray, mask: np.ndarray) -> LatentImage: img = Image.fromarray(image) mask_img = Image.fromarray(mask) latent = vae.masked_encode(img, mask_img)