nodes/space-nuko-ComfyUI-Disco-Diffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

second

Browse Source
This commit is contained in:
space-nuko
2023-05-14 23:20:39 -05:00
parent c63baf9885
commit b76f92b45e
8 changed files with 48 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
diffuse.py
+6 -2
View File
@@ -43,7 +43,7 @@ from .do_run import do_run
# !! }}
#@title Do the Run!
#@markdown `n_batches` ignored with animation modes.
def diffuse(clip, args: DiscoDiffusionSettings, batchNum):
def diffuse(clip, clip_vision, args: DiscoDiffusionSettings, batchNum):
args.display_rate = 20 #@param{type: 'number'}
args.n_batches = 50 #@param{type: 'number'}
@@ -133,6 +133,10 @@ def diffuse(clip, args: DiscoDiffusionSettings, batchNum):
args.seed = seed
args.prompts_series = disco_utils.split_prompts(args.text_prompts, args.max_frames) if args.text_prompts else None,
args.image_prompts_series = disco_utils.split_prompts(args.image_prompts, args.max_frames) if args.image_prompts else None,
args.cut_overview = eval(args.cut_overview)
args.cut_innercut = eval(args.cut_innercut)
args.cut_ic_pow = eval(args.cut_ic_pow)
args.cut_icgray_p = eval(args.cut_icgray_p)
# args = {
# 'batchNum': batchNum,
@@ -267,7 +271,7 @@ def diffuse(clip, args: DiscoDiffusionSettings, batchNum):
gc.collect()
torch.cuda.empty_cache()
try:
do_run(diffusion, model, clip, args, batchNum)
do_run(diffusion, model, clip, clip_vision, args, batchNum)
except KeyboardInterrupt:
pass
finally:
disco_utils.py
+4
View File
@@ -175,6 +175,10 @@ def range_loss(input):
return (input - input.clamp(-1, 1)).pow(2).mean([1, 2, 3])
def alpha_sigma_to_t(alpha, sigma):
return torch.atan2(sigma, alpha) * 2 / math.pi
normalize = T.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])
do_run.py
+27 -18
View File
@@ -22,8 +22,10 @@ import random
from numpy import asarray
from . import py3d_tools as p3dT
from . import disco_xform_utils as dxf
from .CLIP import clip
import comfy.model_management
import comfy.utils
from . import disco_utils
from .make_cutouts import MakeCutouts, MakeCutoutsDango
@@ -90,13 +92,16 @@ def id(x):
return x
def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
def do_run(diffusion, model, clip_, clip_vision, args: DiscoDiffusionSettings, batchNum):
seed = args.seed
print(range(args.start_frame, args.max_frames))
pbar = comfy.utils.ProgressBar(diffusion.num_timesteps - args.skip_steps)
if (args.animation_mode == "3D") and (args.midas_weight > 0.0):
midas_model, midas_transform, midas_net_w, midas_net_h, midas_resize_mode, midas_normalization = init_midas_depth_model(
args.midas_depth_model)
for frame_num in range(args.start_frame, args.max_frames):
if stop_on_next_loop:
break
@@ -104,10 +109,10 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
# display.clear_output(wait=True)
# Print Frame progress if animation mode is on
if args.animation_mode != "None":
batchBar = tqdm(range(args.max_frames), desc="Frames")
batchBar.n = frame_num
batchBar.refresh()
# if args.animation_mode != "None":
# batchBar = tqdm(range(args.max_frames), desc="Frames")
# batchBar.n = frame_num
# batchBar.refresh()
# Inits if not video frames
if args.animation_mode != "Video Input":
@@ -278,7 +283,7 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
print(f'Frame {frame_num} Prompt: {frame_prompt}')
clip_models = [clip] # TODO!!!!!!!!!!!!!!!!!!!!!
clip_models = [clip_] # TODO!!!!!!!!!!!!!!!!!!!!!
model_stats = []
for clip_model in clip_models:
@@ -286,10 +291,13 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
model_stat = {"clip_model": None, "target_embeds": [],
"make_cutouts": None, "weights": []}
model_stat["clip_model"] = clip_model
model_stat["clip_vision_model"] = clip_vision
for prompt in frame_prompt:
prompt = ", ".join(prompt)
txt, weight = disco_utils.parse_prompt(prompt)
txt = clip_model.encode(prompt).float()
# txt = clip_model.encode(prompt).float()
txt = clip_model.encode_text(clip.tokenize(prompt).to(device)).float()
if args.fuzzy_prompt:
for i in range(25):
@@ -310,7 +318,7 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
img, min(args.side_x, args.side_y, *img.size), T.InterpolationMode.LANCZOS)
batch = model_stat["make_cutouts"](TF.to_tensor(
img).to(device).unsqueeze(0).mul(2).sub(1))
embed = clip_model.encode_image(
embed = clip_vision.encode_image(
disco_utils.normalize(batch)).float()
if args.fuzzy_prompt:
for i in range(25):
@@ -337,7 +345,7 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
init = TF.to_tensor(init).to(device).unsqueeze(0).mul(2).sub(1)
if args.perlin_init:
init = disco_utils.regen_perlin()
init = disco_utils.regen_perlin(args.perlin_mode, args.batch_size)
cur_t = None
@@ -385,7 +393,7 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
)
clip_in = disco_utils.normalize(
cuts(x_in.add(1).div(2)))
image_embeds = model_stat["clip_model"].encode_image(
image_embeds = model_stat["clip_vision_model"].encode_image(
clip_in).float()
dists = disco_utils.spherical_dist_loss(image_embeds.unsqueeze(
1), model_stat["target_embeds"].unsqueeze(0))
@@ -421,19 +429,19 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
return grad * magnitude.clamp(max=args.clamp_max) / magnitude
return grad
if args.diffusion_sampling_mode == 'ddim':
if args.MS.diffusion_sampling_mode == 'ddim':
sample_fn = diffusion.ddim_sample_loop_progressive
else:
sample_fn = diffusion.plms_sample_loop_progressive
# image_display = Output()
for i in range(args.n_batches):
if args.animation_mode == 'None':
# if args.animation_mode == 'None':
# display.clear_output(wait=True)
batchBar = tqdm(range(args.n_batches), desc="Batches")
batchBar.n = i
batchBar.refresh()
print('')
# batchBar = tqdm(range(args.n_batches), desc="Batches")
# batchBar.n = i
# batchBar.refresh()
print(f"+++ Batch {i} +++")
# display.display(image_display)
gc.collect()
torch.cuda.empty_cache()
@@ -442,7 +450,7 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
if args.perlin_init:
init = disco_utils.regen_perlin(
args.perlin_mode, args.batch_size)
args.perlin_mode, args.batch_size, True)
symmetry_transformation_fn = id
if args.use_horizontal_symmetry:
@@ -450,7 +458,7 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
if args.use_vertical_symmetry:
symmetry_transformation_fn = vert_symmetry
if args.diffusion_sampling_mode == 'ddim':
if args.MS.diffusion_sampling_mode == 'ddim':
samples = sample_fn(
model,
(args.batch_size, 3, args.side_y, args.side_x),
@@ -482,6 +490,7 @@ def do_run(diffusion, model, clip, args: DiscoDiffusionSettings, batchNum):
# with run_display:
# display.clear_output(wait=True)
for j, sample in enumerate(samples):
pbar.update_absolute(j, len(samples))
cur_t -= 1
intermediateStep = False
if args.steps_per_checkpoint is not None:
make_cutouts.py
+2 -1
View File
@@ -68,7 +68,8 @@ class MakeCutoutsDango(nn.Module):
def __init__(self,
animation_mode: str,
skip_augs,
cut_size, Overview=4,
cut_size,
Overview=4,
InnerCrop = 0, IC_Size_Pow=0.5, IC_Grey_P = 0.2
):
super().__init__()
nodes.py
+7 -3
View File
@@ -1,15 +1,18 @@
import os.path
import comfy.model_management
NODE_FILE = os.path.abspath(__file__)
DISCO_DIFFUSION_ROOT = os.path.dirname(NODE_FILE)
import sys
sys.path.append(os.path.join(DISCO_DIFFUSION_ROOT, "CLIP"))
sys.path.append(os.path.join(DISCO_DIFFUSION_ROOT, "MiDaS"))
sys.path.append(os.path.join(DISCO_DIFFUSION_ROOT, "ResizeRight"))
sys.path.append(os.path.join(DISCO_DIFFUSION_ROOT, "guided-diffusion"))
sys.path.append(os.path.join(DISCO_DIFFUSION_ROOT, "RAFT/core"))
from .CLIP import clip
from .settings import DiscoDiffusionSettings
from .model_settings import ModelSettings
from .diffuse import diffuse
@@ -19,7 +22,6 @@ class DiscoDiffusion:
@classmethod
def INPUT_TYPES(s):
return {"required": {"text": ("STRING", {"multiline": True}),
"clip": ("CLIP", ),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
}}
RETURN_TYPES = ()
@@ -35,8 +37,10 @@ class DiscoDiffusion:
self.settings.setup(self.model_settings)
self.model_settings.setup(self.settings)
def generate(self, text, clip, seed):
diffuse(clip, self.settings, 0)
def generate(self, text, seed):
device = comfy.model_management.get_torch_device()
clip_model = clip.load('ViT-B/32', jit=False)[0].eval().requires_grad_(False).to(device)
diffuse(clip_model, clip_model, self.settings, 0)
return { "ui": { "images": {} } }
requirements.txt
+1
View File
@@ -8,3 +8,4 @@ datetime
pandas
opencv-python
regex
https://github.com/assafshocher/ResizeRight.git
secondary_diffusion_model.py
-4
View File
@@ -25,10 +25,6 @@ def append_dims(x, n):
return x[(Ellipsis, *(None,) * (n - x.ndim))]
def alpha_sigma_to_t(alpha, sigma):
return torch.atan2(sigma, alpha) * 2 / math.pi
def expand_to_planes(x, shape):
return append_dims(x, len(shape)).repeat([1, 1, *shape[2:]])
settings.py
+1 -1
View File
@@ -111,7 +111,7 @@ class DiscoDiffusionSettings:
# @markdown All rotations are provided in degrees.
self.key_frames = True # @param {type:"boolean"}
self.max_frames = 10000 # @param {type:"number"}
self.max_frames = 100 # @param {type:"number"}
# Do not change, currently will not look good. param ['Linear','Quadratic','Cubic']{type:"string"}
self.interp_spline = 'Linear'