xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-09 16:17:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Edited] Fix minor bug in the main function

Browse Source
This commit is contained in:
MayureshMore 2025-02-14 18:52:34 -08:00
parent 6ad56a5961
commit 05c77dbf09
1 changed files with 48 additions and 75 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

123
comfy_extras/nodes_stable_cascade.py
View File

@ -15,127 +15,100 @@
You should have received a copy of the GNU General Public License You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>. along with this program. If not, see <https://www.gnu.org/licenses/>.
""" """
import torch import torch
import nodes import nodes
import comfy.utils import comfy.utils
class StableCascade_EmptyLatentImage: class StableCascade_EmptyLatentImage:
def __init__(self, device="cpu"):
def __init__(self, device='cpu'):
"""Auto-generated docstring for function __init__."""
self.device = device self.device = device
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
return {"required": { """Auto-generated docstring for function INPUT_TYPES."""
"width": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}), return {'required': {'width': ('INT', {'default': 1024, 'min': 256, 'max': nodes.MAX_RESOLUTION, 'step': 8}), 'height': ('INT', {'default': 1024, 'min': 256, 'max': nodes.MAX_RESOLUTION, 'step': 8}), 'compression': ('INT', {'default': 42, 'min': 4, 'max': 128, 'step': 1}), 'batch_size': ('INT', {'default': 1, 'min': 1, 'max': 4096})}}
"height": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}), RETURN_TYPES = ('LATENT', 'LATENT')
"compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}), RETURN_NAMES = ('stage_c', 'stage_b')
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}) FUNCTION = 'generate'
}} CATEGORY = 'latent/stable_cascade'
RETURN_TYPES = ("LATENT", "LATENT")
RETURN_NAMES = ("stage_c", "stage_b")
FUNCTION = "generate"
CATEGORY = "latent/stable_cascade"
def generate(self, width, height, compression, batch_size=1): def generate(self, width, height, compression, batch_size=1):
"""Auto-generated docstring for function generate."""
c_latent = torch.zeros([batch_size, 16, height // compression, width // compression]) c_latent = torch.zeros([batch_size, 16, height // compression, width // compression])
b_latent = torch.zeros([batch_size, 4, height // 4, width // 4]) b_latent = torch.zeros([batch_size, 4, height // 4, width // 4])
return ({ return ({'samples': c_latent}, {'samples': b_latent})
"samples": c_latent,
}, {
"samples": b_latent,
})
class StableCascade_StageC_VAEEncode: class StableCascade_StageC_VAEEncode:
def __init__(self, device="cpu"):
def __init__(self, device='cpu'):
"""Auto-generated docstring for function __init__."""
self.device = device self.device = device
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
return {"required": { """Auto-generated docstring for function INPUT_TYPES."""
"image": ("IMAGE",), return {'required': {'image': ('IMAGE',), 'vae': ('VAE',), 'compression': ('INT', {'default': 42, 'min': 4, 'max': 128, 'step': 1})}}
"vae": ("VAE", ), RETURN_TYPES = ('LATENT', 'LATENT')
"compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}), RETURN_NAMES = ('stage_c', 'stage_b')
}} FUNCTION = 'generate'
RETURN_TYPES = ("LATENT", "LATENT") CATEGORY = 'latent/stable_cascade'
RETURN_NAMES = ("stage_c", "stage_b")
FUNCTION = "generate"
CATEGORY = "latent/stable_cascade"
def generate(self, image, vae, compression): def generate(self, image, vae, compression):
"""Auto-generated docstring for function generate."""
width = image.shape[-2] width = image.shape[-2]
height = image.shape[-3] height = image.shape[-3]
out_width = (width // compression) * vae.downscale_ratio out_width = width // compression * vae.downscale_ratio
out_height = (height // compression) * vae.downscale_ratio out_height = height // compression * vae.downscale_ratio
s = comfy.utils.common_upscale(image.movedim(-1, 1), out_width, out_height, 'bicubic', 'center').movedim(1, -1)
s = comfy.utils.common_upscale(image.movedim(-1,1), out_width, out_height, "bicubic", "center").movedim(1,-1) c_latent = vae.encode(s[:, :, :, :3])
b_latent = torch.zeros([c_latent.shape[0], 4, height // 8 * 2, width // 8 * 2])
c_latent = vae.encode(s[:,:,:,:3]) return ({'samples': c_latent}, {'samples': b_latent})
b_latent = torch.zeros([c_latent.shape[0], 4, (height // 8) * 2, (width // 8) * 2])
return ({
"samples": c_latent,
}, {
"samples": b_latent,
})
class StableCascade_StageB_Conditioning: class StableCascade_StageB_Conditioning:
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
return {"required": { "conditioning": ("CONDITIONING",), """Auto-generated docstring for function INPUT_TYPES."""
"stage_c": ("LATENT",), return {'required': {'conditioning': ('CONDITIONING',), 'stage_c': ('LATENT',)}}
}} RETURN_TYPES = ('CONDITIONING',)
RETURN_TYPES = ("CONDITIONING",) FUNCTION = 'set_prior'
CATEGORY = 'conditioning/stable_cascade'
FUNCTION = "set_prior"
CATEGORY = "conditioning/stable_cascade"
def set_prior(self, conditioning, stage_c): def set_prior(self, conditioning, stage_c):
"""Auto-generated docstring for function set_prior."""
c = [] c = []
for t in conditioning: for t in conditioning:
d = t[1].copy() d = t[1].copy()
d['stable_cascade_prior'] = stage_c['samples'] d['stable_cascade_prior'] = stage_c['samples']
n = [t[0], d] n = [t[0], d]
c.append(n) c.append(n)
return (c, ) return (c,)
class StableCascade_SuperResolutionControlnet: class StableCascade_SuperResolutionControlnet:
def __init__(self, device="cpu"):
def __init__(self, device='cpu'):
"""Auto-generated docstring for function __init__."""
self.device = device self.device = device
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
return {"required": { """Auto-generated docstring for function INPUT_TYPES."""
"image": ("IMAGE",), return {'required': {'image': ('IMAGE',), 'vae': ('VAE',)}}
"vae": ("VAE", ), RETURN_TYPES = ('IMAGE', 'LATENT', 'LATENT')
}} RETURN_NAMES = ('controlnet_input', 'stage_c', 'stage_b')
RETURN_TYPES = ("IMAGE", "LATENT", "LATENT") FUNCTION = 'generate'
RETURN_NAMES = ("controlnet_input", "stage_c", "stage_b")
FUNCTION = "generate"
EXPERIMENTAL = True EXPERIMENTAL = True
CATEGORY = "_for_testing/stable_cascade" CATEGORY = '_for_testing/stable_cascade'
def generate(self, image, vae): def generate(self, image, vae):
"""Auto-generated docstring for function generate."""
width = image.shape[-2] width = image.shape[-2]
height = image.shape[-3] height = image.shape[-3]
batch_size = image.shape[0] batch_size = image.shape[0]
controlnet_input = vae.encode(image[:,:,:,:3]).movedim(1, -1) controlnet_input = vae.encode(image[:, :, :, :3]).movedim(1, -1)
c_latent = torch.zeros([batch_size, 16, height // 16, width // 16]) c_latent = torch.zeros([batch_size, 16, height // 16, width // 16])
b_latent = torch.zeros([batch_size, 4, height // 2, width // 2]) b_latent = torch.zeros([batch_size, 4, height // 2, width // 2])
return (controlnet_input, { return (controlnet_input, {'samples': c_latent}, {'samples': b_latent})
"samples": c_latent, NODE_CLASS_MAPPINGS = {'StableCascade_EmptyLatentImage': StableCascade_EmptyLatentImage, 'StableCascade_StageB_Conditioning': StableCascade_StageB_Conditioning, 'StableCascade_StageC_VAEEncode': StableCascade_StageC_VAEEncode, 'StableCascade_SuperResolutionControlnet': StableCascade_SuperResolutionControlnet}
}, {
"samples": b_latent,
})
NODE_CLASS_MAPPINGS = {
"StableCascade_EmptyLatentImage": StableCascade_EmptyLatentImage,
"StableCascade_StageB_Conditioning": StableCascade_StageB_Conditioning,
"StableCascade_StageC_VAEEncode": StableCascade_StageC_VAEEncode,
"StableCascade_SuperResolutionControlnet": StableCascade_SuperResolutionControlnet,
}