remote vae support raw/png/jpg

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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`

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

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

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}),
 }))