diff --git a/CHANGELOG.md b/CHANGELOG.md index 05b1d9153..973337c64 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,12 +8,13 @@ what can be done? Well, *Huggingface* is now providing *free-of-charge* **remote-VAE-decode** service! - how to use? previous *Full quality* option in UI is replaced with VAE type selector: *Full, Tiny, Remote* currently supports SD15, SDXL and FLUX.1 with more models expected in the near future - availability is limited (log shows '503 Service Unavailable'), - so if remote processing fails SD.Next will fallback to using normal VAE decode process - *note*: only passed item is final latent itself, no user or generate information at all + depending on your bandwidth select mode in *settings -> vae -> raw/png/jpg* + if remote processing fails SD.Next will fallback to using normal VAE decode process + *privacy note*: only passed item is final latent itself without any user or generate information and latent is not stored in the cloud - **UI** - modern ui reorg main tab improve styling, improve scripts/extensions interface and separate ipadapters + - additional ui hints - **Other** - add `--extensions-dir` cli arg and `SD_EXTENSIONSDIR` env variable to specify extensions directory - update `zluda==3.9.0` diff --git a/modules/processing_vae.py b/modules/processing_vae.py index e791f11c6..7dbd9a546 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -230,6 +230,34 @@ def taesd_vae_encode(image): return encoded +def vae_postprocess(tensor, model, output_type='np'): + images = [] + try: + if isinstance(tensor, list) and len(tensor) > 0 and torch.is_tensor(tensor[0]): + tensor = torch.stack(tensor) + if torch.is_tensor(tensor): + if len(tensor.shape) == 3 and tensor.shape[0] == 3: + tensor = tensor.unsqueeze(0) + if hasattr(model, 'video_processor'): + images = model.video_processor.postprocess_video(tensor, output_type='pil') + elif hasattr(model, 'image_processor'): + images = model.image_processor.postprocess(tensor, output_type=output_type) + elif hasattr(model, "vqgan"): + images = tensor.permute(0, 2, 3, 1).cpu().float().numpy() + if output_type == "pil": + images = model.numpy_to_pil(images) + else: + import diffusers + model.image_processor = diffusers.image_processor.VaeImageProcessor() + images = model.image_processor.postprocess(tensor, output_type=output_type) + else: + images = tensor if isinstance(tensor, list) or isinstance(tensor, np.ndarray) else [tensor] + except Exception as e: + shared.log.error(f'VAE postprocess: {e}') + errors.display(e, 'VAE') + return images + + def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, height=None, frames=None): t0 = time.time() model = model or shared.sd_model @@ -242,10 +270,10 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he if vae_type == 'Remote': shared.state.job = 'Remote VAE' from modules.sd_vae_remote import remote_decode - images = remote_decode(latents=latents, width=width, height=height) + tensors = remote_decode(latents=latents, width=width, height=height) shared.state.job = prev_job - if images is not None and len(images) > 0: - return images + if tensors is not None and len(tensors) > 0: + return vae_postprocess(tensors, model, output_type) shared.state.job = 'VAE' if latents.shape[0] == 0: @@ -279,30 +307,13 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he if torch.is_tensor(decoded): decoded = 2.0 * decoded - 1.0 # typical normalized range - if torch.is_tensor(decoded): - if len(decoded.shape) == 3 and decoded.shape[0] == 3: - decoded = decoded.unsqueeze(0) - if hasattr(model, 'video_processor'): - imgs = model.video_processor.postprocess_video(decoded, output_type='pil') - elif hasattr(model, 'image_processor'): - imgs = model.image_processor.postprocess(decoded, output_type=output_type) - elif hasattr(model, "vqgan"): - imgs = decoded.permute(0, 2, 3, 1).cpu().float().numpy() - if output_type == "pil": - imgs = model.numpy_to_pil(imgs) - else: - import diffusers - model.image_processor = diffusers.image_processor.VaeImageProcessor() - imgs = model.image_processor.postprocess(decoded, output_type=output_type) - else: - imgs = decoded if isinstance(decoded, list) or isinstance(decoded, np.ndarray) else [decoded] - + images = vae_postprocess(decoded, model, output_type) shared.state.job = prev_job if shared.cmd_opts.profile or debug: t1 = time.time() shared.log.debug(f'Profile: VAE decode: {t1-t0:.2f}') devices.torch_gc() - return imgs + return images def vae_encode(image, model, vae_type='Full'): # pylint: disable=unused-variable diff --git a/modules/sd_vae_remote.py b/modules/sd_vae_remote.py index 1cd9bb6fe..7d2645b9a 100644 --- a/modules/sd_vae_remote.py +++ b/modules/sd_vae_remote.py @@ -1,6 +1,6 @@ import io import time -import base64 +import json import torch import requests from PIL import Image @@ -8,44 +8,81 @@ from safetensors.torch import _tobytes hf_endpoints = { - 'sd': 'https://lqmfdhmzmy4dw51z.us-east-1.aws.endpoints.huggingface.cloud', - 'sdxl': 'https://m5fxqwyk0r3uu79o.us-east-1.aws.endpoints.huggingface.cloud', - 'f1': 'https://zy1z7fzxpgtltg06.us-east-1.aws.endpoints.huggingface.cloud', + 'sd': 'https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud', + 'sdxl': 'https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud', + 'f1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud', + 'hunyuanvideo': 'https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud', +} +dtypes = { + "float16": torch.float16, + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "uint8": torch.uint8, } def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_type: str = None) -> Image.Image: from modules import devices, shared, errors, modelloader - images = [] + tensors = [] + content = 0 model_type = model_type or shared.sd_model_type url = hf_endpoints.get(model_type, None) if url is None: shared.log.error(f'Decode: type="remote" type={model_type} unsuppported') - return images + return tensors t0 = time.time() modelloader.hf_login() latents = latents.unsqueeze(0) if len(latents.shape) == 3 else latents for i in range(latents.shape[0]): try: latent = latents[i].detach().clone().to(device=devices.cpu, dtype=devices.dtype).unsqueeze(0) - encoded = base64.b64encode(_tobytes(latent, "inputs")).decode("utf-8") - params = {"shape": list(latent.shape), "dtype": str(latent.dtype).split(".", maxsplit=1)[-1]} + params = { + "do_scaling": True, + "input_tensor_type": "binary", + "shape": list(latent.shape), + "dtype": str(latent.dtype).split(".", maxsplit=1)[-1], + } + headers = { "Content-Type": "tensor/binary" } + if shared.opts.remote_vae_type == 'png': + params["image_format"] = "png" + params["output_type"] = "pil" + headers["Accept"] = "image/png" + elif shared.opts.remote_vae_type == 'jpg': + params["image_format"] = "jpg" + params["output_type"] = "pil" + headers["Accept"] = "image/jpeg" + elif shared.opts.remote_vae_type == 'raw': + params["partial_postprocess"] = False + params["output_type"] = "pt" + params["output_tensor_type"] = "binary" + headers["Accept"] = "tensor/binary" if (model_type == 'f1') and (width > 0) and (height > 0): params['width'] = width params['height'] = height response = requests.post( url=url, - json={"inputs": encoded, "parameters": params}, - headers={"Content-Type": "application/json", "Accept": "image/jpeg"}, - timeout=60, + headers=headers, + params=params, + data=_tobytes(latent, "tensor"), + timeout=300, ) if not response.ok: - shared.log.error(f'Decode: type="remote" model={model_type} code={response.status_code} {response.json()}') + shared.log.error(f'Decode: type="remote" model={model_type} code={response.status_code} headers={response.headers} {response.json()}') else: - images.append(Image.open(io.BytesIO(response.content))) + content += len(response.content) + if shared.opts.remote_vae_type == 'raw': + shape = json.loads(response.headers["shape"]) + dtype = response.headers["dtype"] + tensor = torch.frombuffer(bytearray(response.content), dtype=dtypes[dtype]).reshape(shape) + tensors.append(tensor) + elif shared.opts.remote_vae_type == 'jpg' or shared.opts.remote_vae_type == 'png': + image = Image.open(io.BytesIO(response.content)).convert("RGB") + tensors.append(image) except Exception as e: shared.log.error(f'Decode: type="remote" model={model_type} {e}') errors.display(e, 'VAE') + if len(tensors) > 0 and shared.opts.remote_vae_type == 'raw': + tensors = torch.cat(tensors, dim=0) t1 = time.time() - shared.log.debug(f'Decode: type="remote" model={model_type} args={params} images={images} time={t1-t0:.3f}s') - return images + shared.log.debug(f'Decode: type="remote" model={model_type} mode={shared.opts.remote_vae_type} args={params} bytes={content} time={t1-t0:.3f}s') + return tensors diff --git a/modules/shared.py b/modules/shared.py index a55355c88..962bf05d0 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -423,6 +423,7 @@ options_templates.update(options_section(('vae_encoder', "Variable Auto Encoder" "diffusers_vae_tile_overlap": OptionInfo(0.25, "VAE tile overlap", gr.Slider, {"minimum": 0, "maximum": 0.95, "step": 0.05 }), "sd_vae_sliced_encode": OptionInfo(False, "VAE sliced encode", gr.Checkbox, {"visible": not native}), "nan_skip": OptionInfo(False, "Skip Generation if NaN found in latents", gr.Checkbox), + "remote_vae_type": OptionInfo('raw', "Remote VAE image type", gr.Dropdown, {"choices": ['raw', 'jpg', 'png']}), "rollback_vae": OptionInfo(False, "Attempt VAE roll back for NaN values", gr.Checkbox, {"visible": not native}), }))