adding custom auto encoder

app.py

@@ -135,6 +135,19 @@ def requires_safety_checker_change(requires_safety_checker, config):
 
     return config, str(config), assemble_code(config)
 
+def auto_encoders_change(auto_encoder, config):
+    
+    if str(auto_encoder) != 'None' and type(auto_encoder) != list:
+        
+        auto_encoder_description = auto_encoders[auto_encoder]
+        
+    else:
+        auto_encoder_description = ''
+
+    config = set_config(config, 'auto_encoder', auto_encoder)
+
+    return auto_encoder_description, config, str(config), assemble_code(config)
+
 def schedulers_change(scheduler, config):
     
     if str(scheduler) != 'None' and type(scheduler) != list:
@@ -222,6 +235,7 @@ appConfig = load_app_config()
 models = appConfig.get("models", {})
 schedulers = appConfig.get("schedulers", {})
 devices =  appConfig.get("devices", [])
+auto_encoders = appConfig.get("auto_encoders", [])
 
 # interface
 with gr.Blocks(analytics_enabled=False) as demo:
@@ -246,7 +260,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
         out_model_description = gr.Textbox(value="", label="Description")
     with gr.Row():
         with gr.Column(scale=1):
-            in_trigger_token = gr.Textbox(value=config.value["trigger_token"], label="Trigger Token")
+            in_trigger_token = gr.Textbox(value=config.value["trigger_token"], label="Trigger Token", info="will be added to your prompt to `activate` a fine tuned model")
             in_use_safetensors = gr.Radio(label="Use safe tensors:", choices=["True", "False"], interactive=False)
             in_model_refiner = gr.Dropdown(value=config.value["refiner"], choices=["none"], label="Refiner", allow_custom_value=True, multiselect=False)
         with gr.Column(scale=1):
@@ -271,18 +285,23 @@ with gr.Blocks(analytics_enabled=False) as demo:
         in_inference_steps = gr.Number(label="Inference steps", value=config.value["inference_steps"], info="Each step improves the final result but also results in higher computation")
         in_manual_seed = gr.Number(label="Manual seed", value=config.value["manual_seed"], info="Set this to -1 or leave it empty to randomly generate an image. A fixed value will result in a similar image for every run")
         in_guidance_scale = gr.Slider(minimum=0, maximum=100, step=0.1, label="Guidance Scale", value=config.value["guidance_scale"], info="A low guidance scale leads to a faster inference time, with the drawback that negative prompts don’t have any effect on the denoising process.")
+    with gr.Row():
+        gr.Markdown("### Auto Encoder")
     with gr.Row():
         gr.Markdown("**VAE** stands for Variational Autoencoders. An 'autoencoder' is an artificial neural network that is able to encode input data and decode to output data to bascially recreate the input. The VAE whereas adds a couple of additional layers of complexity to create new and unique output.")
     with gr.Row():
-        in_enable_vae_slicing = gr.Radio(label="Enable VAE slicing:", value=config.value["enable_vae_slicing"], choices=["True", "False"], info="decoding the batches of latents one image at a time, which may reduce memory usage, see https://huggingface.co/docs/diffusers/main/en/optimization/memory")
-        in_enable_vae_tiling= gr.Radio(label="Enable VAE tiling:", value=config.value["enable_vae_tiling"], choices=["True", "False"], info="splitting the image into overlapping tiles, decoding the tiles, and then blending the outputs together to compose the final image, see https://huggingface.co/docs/diffusers/main/en/optimization/memory")
+        with gr.Column():
+            in_auto_encoders = gr.Dropdown(value="None", choices=list(auto_encoders.keys()), label="Auto encoder", info="leave empty to not add an auto encoder")
+            out_auto_encoder_description = gr.Textbox(value="", label="Description")
+        with gr.Column():
+            in_enable_vae_slicing = gr.Radio(label="Enable VAE slicing:", value=config.value["enable_vae_slicing"], choices=["True", "False"], info="decoding the batches of latents one image at a time, which may reduce memory usage, see https://huggingface.co/docs/diffusers/main/en/optimization/memory")
+            in_enable_vae_tiling= gr.Radio(label="Enable VAE tiling:", value=config.value["enable_vae_tiling"], choices=["True", "False"], info="splitting the image into overlapping tiles, decoding the tiles, and then blending the outputs together to compose the final image, see https://huggingface.co/docs/diffusers/main/en/optimization/memory")
         
     gr.Markdown("### Output")
     with gr.Row():
         btn_start_pipeline = gr.Button(value="Run", variant="primary")
         btn_stop_pipeline = gr.Button(value="Stop", variant="stop")
     with gr.Row():
-        # out_result = gr.Textbox(label="Status", value="")
         out_image = gr.Image()
         out_code = gr.Code(assemble_code(config.value), label="Code")
     with gr.Row():
@@ -303,6 +322,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
     in_inference_steps.change(inference_steps_change, inputs=[in_inference_steps, config], outputs=[config, out_config, out_code])
     in_manual_seed.change(manual_seed_change, inputs=[in_manual_seed, config], outputs=[config, out_config, out_code])
     in_guidance_scale.change(guidance_scale_change, inputs=[in_guidance_scale, config], outputs=[config, out_config, out_code])
+    in_auto_encoders.change(auto_encoders_change, inputs=[in_auto_encoders, config], outputs=[out_auto_encoder_description, config, out_config, out_code])
     in_enable_vae_slicing.change(enable_vae_slicing_change, inputs=[in_enable_vae_slicing, config], outputs=[config, out_config, out_code])
     in_enable_vae_tiling.change(enable_vae_tiling_change, inputs=[in_enable_vae_tiling, config], outputs=[config, out_config, out_code])
     in_prompt.change(prompt_change, inputs=[in_prompt, config], outputs=[config, out_config, out_code])
@@ -325,6 +345,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
             in_variant,
             in_safety_checker,
             in_requires_safety_checker,
+            in_auto_encoders,
             in_enable_vae_slicing,
             in_enable_vae_tiling,
             in_schedulers,

appConfig.json

@@ -72,6 +72,12 @@
     "backup_devices": [
         "cpu", "cuda", "ipu", "xpu", "mkldnn", "opengl", "opencl", "ideep", "hip", "ve", "fpga", "ort", "xla", "lazy", "vulkan", "mps", "meta", "hpu", "mtia", "privateuseone", "gpu"
     ],
+    "auto_encoders": {
+        "None": "",
+        "stabilityai/sdxl-vae": "finetuned auto encoder for stable diffusion models, see https://huggingface.co/stabilityai/sdxl-vae",
+        "madebyollin/sdxl-vae-fp16-fix": "stable diffusion models encoder with fp16 precision, see https://huggingface.co/madebyollin/sdxl-vae-fp16-fix",
+        "stabilityai/sd-vae-ft-mse": "works best with CompVis/stable-diffusion-v1-4, see https://huggingface.co/stabilityai/sd-vae-ft-mse"
+    },
     "schedulers": {
         "DDPMScheduler":                    "Denoising Diffusion Probabilistic Model",  
         "DDIMScheduler":                    "Denoising Diffusion Incremental Sampling, efficient image generation, might require more tunin",  

config.py

@@ -46,6 +46,7 @@ def get_initial_config():
         "refiner": "none",
         "safety_checker": "False",
         "requires_safety_checker": "False",
+        "auto_encoder": None,
         "enable_vae_slicing": "True",
         "enable_vae_tiling": "True",
         "manual_seed": 42,
@@ -88,6 +89,7 @@ def get_config_from_url(initial_config, request: Request):
             return_config['variant'],
             return_config['safety_checker'],
             return_config['requires_safety_checker'],
+            return_config['auto_encoder'],
             return_config['enable_vae_slicing'],
             return_config['enable_vae_tiling'],
             return_config['scheduler'],
@@ -151,6 +153,9 @@ def assemble_code(str_config):
             torch_dtype=data_type, 
             variant=variant).to(device)''')
     
+    if str(config["auto_encoder"]).lower() != 'none': 
+        code.append(f'pipeline.vae = AutoencoderKL.from_pretrained("{config["auto_encoder"]}", torch_dtype=data_type).to(device)')
+    
     if str(config["cpu_offload"]).lower() != 'false': code.append("pipeline.enable_model_cpu_offload()")
     if str(config["enable_vae_slicing"]).lower() != 'false': code.append("pipeline.enable_vae_slicing()")
     if str(config["enable_vae_tiling"]).lower() != 'false': code.append("pipeline.enable_vae_tiling()")