further Fooocus import

2023-08-16 13:56:21 +02:00
parent c9fba811a5
commit ad2e23b150
5 changed files with 227 additions and 70 deletions
@@ -1,11 +1,18 @@
+import os
 import torch
 import comfy.model_management
+import comfy.utils
 import math
 import numpy as np
 import latent_preview
 from comfy.sample import *

 from . import samplers_advanced
+from .adm_patch import patch_negative_adm
+
+
+patch_negative_adm()
+

 def sample_refined(model, refiner_model, noise, steps, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
    """
@@ -82,6 +89,85 @@ def ksampler_refined(model, refiner_model, seed, steps, cfg, sampler_name, sched
    out["samples"] = samples
    return (out, )

+
+@torch.no_grad()
+def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive, refiner_negative, latent,
+                          seed=None, steps=30, refiner_switch_step=20, cfg=7.0, sampler_name='dpmpp_2m_sde_gpu',
+                          scheduler='karras', denoise=1.0, disable_noise=False, start_step=None, last_step=None,
+                          force_full_denoise=False, callback_function=None):
+    # SCHEDULERS = ["normal", "karras", "exponential", "simple", "ddim_uniform"]
+    # SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",
+    #             "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
+    #             "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "ddim", "uni_pc", "uni_pc_bh2"]
+
+    seed = seed if isinstance(seed, int) else random.randint(1, 2 ** 64)
+
+    device = comfy.model_management.get_torch_device()
+    latent_image = latent["samples"]
+
+    if disable_noise:
+        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+    else:
+        batch_inds = latent["batch_index"] if "batch_index" in latent else None
+        noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
+
+    noise_mask = None
+    if "noise_mask" in latent:
+        noise_mask = latent["noise_mask"]
+
+    preview_format = "JPEG"
+    if preview_format not in ["JPEG", "PNG"]:
+        preview_format = "JPEG"
+
+    previewer = latent_preview.get_previewer(device, model.model.latent_format)
+
+    pbar = comfy.utils.ProgressBar(steps)
+    def callback(step, x0, x, total_steps):
+        preview_bytes = None
+        if previewer:
+            preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
+        pbar.update_absolute(step + 1, total_steps, preview_bytes)
+
+    sigmas = None
+    disable_pbar = False
+
+    if noise_mask is not None:
+        noise_mask = prepare_mask(noise_mask, noise.shape, device)
+
+    comfy.model_management.load_model_gpu(model)
+
+    noise = noise.to(device)
+    latent_image = latent_image.to(device)
+
+    positive_copy = broadcast_cond(positive, noise.shape[0], device)
+    negative_copy = broadcast_cond(negative, noise.shape[0], device)
+
+    refiner_positive_copy = broadcast_cond(refiner_positive, noise.shape[0], device)
+    refiner_negative_copy = broadcast_cond(refiner_negative, noise.shape[0], device)
+
+    models = load_additional_models(positive, negative, model.model_dtype())
+
+    sampler = samplers_advanced.KSamplerWithRefiner(model=model, refiner_model=refiner, steps=steps, device=device,
+                                  sampler=sampler_name, scheduler=scheduler,
+                                  denoise=denoise, model_options=model.model_options)
+
+    samples = sampler.sample(noise, positive_copy, negative_copy, refiner_positive=refiner_positive_copy,
+                             refiner_negative=refiner_negative_copy, refiner_switch_step=refiner_switch_step,
+                             cfg=cfg, latent_image=latent_image,
+                             start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise,
+                             denoise_mask=noise_mask, sigmas=sigmas, callback_function=callback, disable_pbar=disable_pbar,
+                             seed=seed)
+
+    samples = samples.cpu()
+
+    cleanup_additional_models(models)
+
+    out = latent.copy()
+    out["samples"] = samples
+
+    return out
+
+
 class KSamplerWithRefiner:
    @classmethod
    def INPUT_TYPES(s):
@@ -90,10 +176,11 @@ class KSamplerWithRefiner:
                    "refiner_model": ("MODEL",),
                    "add_noise": (["enable", "disable"], ),
                    "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
-                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                    "steps": ("INT", {"default": 30, "min": 1, "max": 10000}),
+                    "refiner_switch_step": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                    "cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0}),
+                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "dpmpp_2m_sde_gpu", }),
+                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "karras", }),
                    "positive": ("CONDITIONING", ),
                    "negative": ("CONDITIONING", ),
                    "refiner_positive": ("CONDITIONING", ),
@@ -110,11 +197,12 @@ class KSamplerWithRefiner:

    CATEGORY = "sampling"

-    def sample(self, model, refiner_model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0):
+    def sample(self, model, refiner_model, add_noise, noise_seed, steps, refiner_switch_step, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0):
        force_full_denoise = True
        if return_with_leftover_noise == "enable":
            force_full_denoise = False
        disable_noise = False
        if add_noise == "disable":
            disable_noise = True
-        return ksampler_refined(model, refiner_model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
+        # return ksampler_refined(model, refiner_model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
+        return ksampler_with_refiner(model, positive, negative, refiner_model, refiner_positive, refiner_negative, latent_image, noise_seed, steps, refiner_switch_step, cfg, sampler_name, scheduler, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
@@ -0,0 +1,33 @@
+import torch
+import comfy.model_base
+
+
+def sdxl_encode_adm_patched(self, **kwargs):
+    clip_pooled = kwargs["pooled_output"]
+    width = kwargs.get("width", 768)
+    height = kwargs.get("height", 768)
+    crop_w = kwargs.get("crop_w", 0)
+    crop_h = kwargs.get("crop_h", 0)
+    target_width = kwargs.get("target_width", width)
+    target_height = kwargs.get("target_height", height)
+
+    if kwargs.get("prompt_type", "") == "negative":
+        width *= 0.8
+        height *= 0.8
+    elif kwargs.get("prompt_type", "") == "positive":
+        width *= 1.5
+        height *= 1.5
+
+    out = []
+    out.append(self.embedder(torch.Tensor([height])))
+    out.append(self.embedder(torch.Tensor([width])))
+    out.append(self.embedder(torch.Tensor([crop_h])))
+    out.append(self.embedder(torch.Tensor([crop_w])))
+    out.append(self.embedder(torch.Tensor([target_height])))
+    out.append(self.embedder(torch.Tensor([target_width])))
+    flat = torch.flatten(torch.cat(out))[None, ]
+    return torch.cat((clip_pooled.to(flat.device), flat), dim=1)
+
+
+def patch_negative_adm():
+    comfy.model_base.SDXL.encode_adm = sdxl_encode_adm_patched
@@ -0,0 +1,35 @@
+import threading
+import cv2
+
+
+buffer = []
+
+
+def worker():
+    global buffer
+    while True:
+        cv2.waitKey(50)
+        try:
+            if len(buffer) > 0:
+                task = buffer.pop(0)
+                if task is None:
+                    cv2.destroyAllWindows()
+                else:
+                    flag, img, title = task
+                    cv2.imshow(flag, img)
+                    cv2.setWindowTitle(flag, title)
+                    cv2.setWindowProperty(flag, cv2.WND_PROP_TOPMOST, 1)
+        except Exception as e:
+            print(e)
+    pass
+
+
+def show_preview(flag, img, title='preview'):
+    buffer.append((flag, img[..., ::-1].copy(), title))
+
+
+def close_all_preview():
+    buffer.append(None)
+
+
+threading.Thread(target=worker, daemon=True).start()
@@ -19,8 +19,8 @@ class KSamplerWithRefiner:
        self.model_patcher = model
        self.refiner_model_patcher = refiner_model

-        self.model = model
-        self.refiner_model = refiner_model
+        self.model = model.model
+        self.refiner_model = refiner_model.model

        self.model_denoise = CFGNoisePredictor(self.model)
        self.refiner_model_denoise = CFGNoisePredictor(self.refiner_model)
@@ -271,32 +271,6 @@
        "sdxl_vae.safetensors"
      ]
    },
-    {
-      "id": 28,
-      "type": "SaveImage",
-      "pos": [
-        2465,
-        298
-      ],
-      "size": {
-        "0": 315,
-        "1": 58
-      },
-      "flags": {},
-      "order": 17,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "images",
-          "type": "IMAGE",
-          "link": 51
-        }
-      ],
-      "properties": {},
-      "widgets_values": [
-        "ComfyUI"
-      ]
-    },
    {
      "id": 17,
      "type": "CLIPTextEncodeSDXL",
@@ -584,7 +558,7 @@
      ],
      "size": {
        "0": 315,
-        "1": 334
+        "1": 546
      },
      "flags": {},
      "order": 11,
@@ -629,12 +603,12 @@
        "enable",
        1,
        "fixed",
-        20,
+        30,
        8,
-        "dpmpp_2m",
+        "dpmpp_2m_sde_gpu",
        "karras",
        0,
-        16,
+        20,
        "enable"
      ]
    },
@@ -707,9 +681,10 @@
        "enable",
        1,
        "fixed",
+        30,
        20,
-        8,
-        "dpmpp_2m",
+        7,
+        "dpmpp_2m_sde_gpu",
        "karras",
        0,
        10000,
@@ -725,7 +700,7 @@
      ],
      "size": {
        "0": 315,
-        "1": 334
+        "1": 546
      },
      "flags": {},
      "order": 13,
@@ -770,11 +745,11 @@
        "disable",
        1,
        "fixed",
-        20,
+        30,
        8,
-        "dpmpp_2m",
+        "dpmpp_2m_sde_gpu",
        "karras",
-        16,
+        20,
        10000,
        "disable"
      ]
@@ -819,32 +794,6 @@
        "Node name for S&R": "VAEDecode"
      }
    },
-    {
-      "id": 9,
-      "type": "SaveImage",
-      "pos": [
-        2014.1724880000002,
-        -373.70633999999995
-      ],
-      "size": {
-        "0": 210,
-        "1": 58
-      },
-      "flags": {},
-      "order": 16,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "images",
-          "type": "IMAGE",
-          "link": 9
-        }
-      ],
-      "properties": {},
-      "widgets_values": [
-        "ComfyUI"
-      ]
-    },
    {
      "id": 27,
      "type": "VAEDecode",
@@ -885,6 +834,58 @@
      "properties": {
        "Node name for S&R": "VAEDecode"
      }
+    },
+    {
+      "id": 9,
+      "type": "SaveImage",
+      "pos": [
+        2014.1724880000002,
+        -373.70633999999995
+      ],
+      "size": {
+        "0": 210,
+        "1": 270
+      },
+      "flags": {},
+      "order": 16,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 9
+        }
+      ],
+      "properties": {},
+      "widgets_values": [
+        "ComfyUI"
+      ]
+    },
+    {
+      "id": 28,
+      "type": "SaveImage",
+      "pos": [
+        2465,
+        298
+      ],
+      "size": {
+        "0": 315,
+        "1": 270
+      },
+      "flags": {},
+      "order": 17,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 51
+        }
+      ],
+      "properties": {},
+      "widgets_values": [
+        "ComfyUI"
+      ]
    }
  ],
  "links": [