remote vae encode

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

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### TODO
   - Gemma3 requires `git+https://github.com/huggingface/transformers@v4.49.0-Gemma-3`
+  - Remote VAE encode for SD15 and Flux.1: <https://github.com/huggingface/diffusers/issues/11069>
 
 - **Models**  
   - [THUDM CogView 4 6B](https://huggingface.co/THUDM/CogView4-6B)  
@@ -15,6 +16,10 @@
     download text encoders into folder set in settings -> system paths -> text encoders (default is `models/Text-encoder`)  
     load using *settings -> text encoder*
     *tip*: add *sd_text_encoder* to your *settings -> user interface -> quicksettings* list to have it appear at the top of the ui  
+- **Remote VAE**
+  - add support for remote vae encode in addition to remote vae decode  
+  - used by *img2img, inpaint, hires, detailer*  
+  - remote vae encode is disabled by default, you can enable it in *settings -> variable auto-encoder*  
 - **Caption/VLM**  
   - [Google Gemma 3 4B](https://huggingface.co/google/gemma-3-4b-it)  
     simply select from list of available models in caption tab  

modules/interrogate/vqa.py

@@ -150,6 +150,9 @@ def qwen(question: str, image: Image.Image, repo: str = None, system_prompt: str
 
 def gemma(question: str, image: Image.Image, repo: str = None, system_prompt: str = None):
     global processor, model, loaded # pylint: disable=global-statement
+    if not hasattr(transformers, 'Gemma3ForConditionalGeneration'):
+        shared.log.error(f'Interrogate: vlm="{repo}" gemma is not available')
+        return ''
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
         model = transformers.Gemma3ForConditionalGeneration.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
@@ -521,9 +524,7 @@ def interrogate(question, system_prompt, prompt, image, model_name, quiet:bool=F
     if shared.opts.interrogate_offload and model is not None:
         model.to(devices.cpu)
     devices.torch_gc()
-    print('HERE1', answer)
     answer = clean(answer, question)
-    print('HERE2', answer)
     t1 = time.time()
     if not quiet:
         shared.log.debug(f'Interrogate: type=vlm model="{model_name}" repo="{vqa_model}" args={get_kwargs()} time={t1-t0:.2f}')

modules/postprocess/yolo.py

@@ -265,6 +265,7 @@ class YoloRestorer(Detailer):
                 'inpaint_full_res_padding': shared.opts.detailer_padding,
                 'width': resolution,
                 'height': resolution,
+                'vae_type': orig_p.get('vae_type', 'Full'),
             }
             if args['denoising_strength'] == 0:
                 shared.log.debug(f'Detailer: model="{name}" strength=0 skip')

modules/processing_args.py

@@ -38,13 +38,20 @@ def task_specific_kwargs(p, model):
             model.register_to_config(requires_aesthetics_score = False)
         if 'hires' not in p.ops:
             p.ops.append('img2img')
+        if p.vae_type == 'Remote':
+            from modules.sd_vae_remote import remote_encode
+            p.init_images = remote_encode(p.init_images)
         task_args = {
             'image': p.init_images,
             'strength': p.denoising_strength,
         }
         if model.__class__.__name__ == 'FluxImg2ImgPipeline': # needs explicit width/height
-            p.width = 8 * math.ceil(p.init_images[0].width / 8)
-            p.height = 8 * math.ceil(p.init_images[0].height / 8)
+            if torch.is_tensor(p.init_images[0]):
+                p.width = p.init_images[0].shape[-1] * 16
+                p.height = p.init_images[0].shape[-2] * 16
+            else:
+                p.width = 8 * math.ceil(p.init_images[0].width / 8)
+                p.height = 8 * math.ceil(p.init_images[0].height / 8)
             task_args['width'], task_args['height'] = p.width, p.height
         if model.__class__.__name__ == 'OmniGenPipeline':
             p.width = 16 * math.ceil(p.init_images[0].width / 16)
@@ -70,9 +77,14 @@ def task_specific_kwargs(p, model):
         else:
             p.ops.append('inpaint')
         width, height = processing_helpers.resize_init_images(p)
+        mask_image = p.task_args.get('image_mask', None) or getattr(p, 'image_mask', None) or getattr(p, 'mask', None)
+        if p.vae_type == 'Remote':
+            from modules.sd_vae_remote import remote_encode
+            p.init_images = remote_encode(p.init_images)
+            # mask_image = remote_encode(mask_image)
         task_args = {
             'image': p.init_images,
-            'mask_image': p.task_args.get('image_mask', None) or getattr(p, 'image_mask', None) or getattr(p, 'mask', None),
+            'mask_image': mask_image,
             'strength': p.denoising_strength,
             'height': height,
             'width': width,

modules/sd_vae_remote.py

@@ -1,3 +1,4 @@
+from typing import List
 import io
 import time
 import json
@@ -38,7 +39,6 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
     t0 = time.time()
     modelloader.hf_login()
     latents = latents.unsqueeze(0) if len(latents.shape) == 3 else latents
-    from diffusers.utils.remote_utils import remote_decode
 
     for i in range(latents.shape[0]):
         try:
@@ -97,25 +97,40 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
     return tensors
 
 
-def remote_encode(image: Image.Image, model_type: str = None) -> torch.Tensor:
+def remote_encode(images: List[Image.Image], model_type: str = None) -> torch.Tensor:
+    from diffusers.utils import remote_utils
     from modules import devices, shared, errors, modelloader
+    if not shared.opts.remote_vae_encode:
+        return images
     tensors = []
     model_type = model_type or shared.sd_model_type
-    url = hf_decode_endpoints.get(model_type, None)
+    url = hf_encode_endpoints.get(model_type, None)
     if url is None:
         shared.log.error(f'Decode: type="remote" type={model_type} unsuppported')
-        return tensors
+        return images
     t0 = time.time()
     modelloader.hf_login()
 
-    try:
-        params = {}
-        content = 0
-        tensor = None
-    except Exception as e:
-        shared.log.error(f'Encode: type="remote" model={model_type} {e}')
-        errors.display(e, 'VAE')
+    if isinstance(images, Image.Image):
+        images = [images]
+    for init_image in images:
+        try:
+            init_latent = remote_utils.remote_encode(
+                endpoint=url,
+                image=init_image,
+                scaling_factor = shared.sd_model.vae.config.get("scaling_factor", None),
+                shift_factor = shared.sd_model.vae.config.get("shift_factor", None),
+            )
+            tensors.append(init_latent)
+        except Exception as e:
+            shared.log.error(f'Encode: type="remote" model={model_type} {e}')
+            errors.display(e, 'VAE')
 
+    if len(tensors) > 0 and torch.is_tensor(tensors[0]):
+        tensors = torch.cat(tensors, dim=0)
+        tensors = tensors.to(dtype=devices.dtype)
+    else:
+        return images
     t1 = time.time()
-    shared.log.debug(f'Encode: type="remote" model={model_type} mode={shared.opts.remote_vae_type} args={params} image={image} bytes={content} time={t1-t0:.3f}s')
-    return tensor
+    shared.log.debug(f'Encode: type="remote" model={model_type} mode={shared.opts.remote_vae_type} image={images} latent={tensors.shape} time={t1-t0:.3f}s')
+    return tensors

modules/shared.py

@@ -424,6 +424,7 @@ options_templates.update(options_section(('vae_encoder', "Variable Auto Encoder"
     "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']}),
+    "remote_vae_encode": OptionInfo(False, "Remote VAE for encode"),
     "rollback_vae": OptionInfo(False, "Attempt VAE roll back for NaN values", gr.Checkbox, {"visible": not native}),
 }))