75 lines
2.5 KiB
Python
75 lines
2.5 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
|
||
|
|
|
||
|
|
|
||
|
|
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": 32, "max": 64, "step": 1}),
|
||
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64})
|
||
|
|
}}
|
||
|
|
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_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,
|
||
|
|
}
|