110 lines
3.7 KiB
Python
110 lines
3.7 KiB
Python
"""
|
|
This file is part of ComfyUI.
|
|
Copyright (C) 2024 Stability AI
|
|
|
|
This program is free software: you can redistribute it and/or modify
|
|
it under the terms of the GNU General Public License as published by
|
|
the Free Software Foundation, either version 3 of the License, or
|
|
(at your option) any later version.
|
|
|
|
This program is distributed in the hope that it will be useful,
|
|
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
GNU General Public License for more details.
|
|
|
|
You should have received a copy of the GNU General Public License
|
|
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
"""
|
|
|
|
import torch
|
|
import nodes
|
|
import comfy.utils
|
|
|
|
|
|
class StableCascade_EmptyLatentImage:
|
|
def __init__(self, device="cpu"):
|
|
self.device = device
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
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})
|
|
}}
|
|
RETURN_TYPES = ("LATENT", "LATENT")
|
|
RETURN_NAMES = ("stage_c", "stage_b")
|
|
FUNCTION = "generate"
|
|
|
|
CATEGORY = "_for_testing/stable_cascade"
|
|
|
|
def generate(self, width, height, compression, batch_size=1):
|
|
c_latent = torch.zeros([batch_size, 16, height // compression, width // compression])
|
|
b_latent = torch.zeros([batch_size, 4, height // 4, width // 4])
|
|
return ({
|
|
"samples": c_latent,
|
|
}, {
|
|
"samples": b_latent,
|
|
})
|
|
|
|
class StableCascade_StageC_VAEEncode:
|
|
def __init__(self, device="cpu"):
|
|
self.device = device
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"image": ("IMAGE",),
|
|
"vae": ("VAE", ),
|
|
"compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}),
|
|
}}
|
|
RETURN_TYPES = ("LATENT", "LATENT")
|
|
RETURN_NAMES = ("stage_c", "stage_b")
|
|
FUNCTION = "generate"
|
|
|
|
CATEGORY = "_for_testing/stable_cascade"
|
|
|
|
def generate(self, image, vae, compression):
|
|
width = image.shape[-2]
|
|
height = image.shape[-3]
|
|
out_width = (width // 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)
|
|
|
|
c_latent = vae.encode(s[:,:,:,:3])
|
|
b_latent = torch.zeros([c_latent.shape[0], 4, height // 4, width // 4])
|
|
return ({
|
|
"samples": c_latent,
|
|
}, {
|
|
"samples": b_latent,
|
|
})
|
|
|
|
class StableCascade_StageB_Conditioning:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "conditioning": ("CONDITIONING",),
|
|
"stage_c": ("LATENT",),
|
|
}}
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
|
|
FUNCTION = "set_prior"
|
|
|
|
CATEGORY = "_for_testing/stable_cascade"
|
|
|
|
def set_prior(self, conditioning, stage_c):
|
|
c = []
|
|
for t in conditioning:
|
|
d = t[1].copy()
|
|
d['stable_cascade_prior'] = stage_c['samples']
|
|
n = [t[0], d]
|
|
c.append(n)
|
|
return (c, )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"StableCascade_EmptyLatentImage": StableCascade_EmptyLatentImage,
|
|
"StableCascade_StageB_Conditioning": StableCascade_StageB_Conditioning,
|
|
"StableCascade_StageC_VAEEncode": StableCascade_StageC_VAEEncode,
|
|
}
|