57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
import io
|
|
import base64
|
|
import time
|
|
|
|
import torch.nn.functional as tnnf
|
|
import comfy.utils
|
|
from server import PromptServer
|
|
from PIL import Image
|
|
|
|
def scale_latent(latent, max_pixels=1000000):
|
|
h = latent.shape[2]
|
|
w = latent.shape[3]
|
|
pixels = h * 8 * w * 8
|
|
if pixels <= max_pixels:
|
|
return latent
|
|
scale = (max_pixels / pixels) ** 0.5
|
|
new_h = int(h * scale)
|
|
new_w = int(w * scale)
|
|
new_h = max(1,new_h)
|
|
new_w = max(1,new_w)
|
|
return tnnf.interpolate(latent, size = (new_h, new_w), mode = "bilinear",
|
|
align_corners = False)
|
|
|
|
def tensor_to_base64(image):
|
|
# Receive a vae image and change it to base64, then send it to front-end.
|
|
image = image.cpu().clamp(0, 1)
|
|
image = (image[0].numpy() * 255).astype("uint8")
|
|
img = Image.fromarray(image)
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format = "PNG")
|
|
return base64.b64encode(buffer.getvalue()).decode()
|
|
|
|
def create_preview_callback(model, vae, steps, preview_decode = "Scale", interval = 1, sync_preview = True):
|
|
# Sampler's callback, it has two mode.
|
|
# Sync: The HD preview will sync with latent's step, one step = one preview image.
|
|
# Fast: The HD preview will be rendered once at intervals.
|
|
pbar = comfy.utils.ProgressBar(steps)
|
|
def inner_callback(step, x0, x, total_steps):
|
|
if(steps % interval == 0 or steps == total_steps):
|
|
latent = (model.model.latent_format.process_out(x0))
|
|
if preview_decode == "Scale":
|
|
latent = scale_latent(latent)
|
|
image = vae.decode(latent)
|
|
data = tensor_to_base64(image)
|
|
PromptServer.instance.send_sync(
|
|
"high_preview", {
|
|
"image": data,
|
|
"step": step + 1,
|
|
"total":total_steps
|
|
}
|
|
)
|
|
if sync_preview:
|
|
# Give browser some time to flash image.
|
|
time.sleep(0.15)
|
|
# Not use default preview window, so there is None.
|
|
pbar.update_absolute(step + 1, total_steps, None)
|
|
return inner_callback |