Refactor vae assignment in handlers.py and load_models.py

tabs/images/handlers.py

@@ -68,7 +68,7 @@ def get_pipe(request: BaseReq | BaseImg2ImgReq | BaseInpaintReq):
             if request.vae:
                 pipe_args["pipeline"].vae = sdxl_vae if model['loader'] == 'sdxl' else flux_vae
             elif not request.vae:
-                pipe_args["pipeline"].vae = None
+                pipe_args["pipeline"].vae = None if model['loader'] == 'sdxl' else flux_vae
             
             # Set Scheduler
             pipe_args["pipeline"].scheduler = get_scheduler(pipe_args["pipeline"], request.scheduler)

tabs/images/load_models.py

@@ -18,7 +18,6 @@ def init_sys():
         try:
             model['pipeline'] = AutoPipelineForText2Image.from_pretrained(
                 model['repo_id'],
-                vae=None,
                 torch_dtype=model['compute_type'],
                 safety_checker=None,
                 variant="fp16"
@@ -26,7 +25,6 @@ def init_sys():
         except:
             model['pipeline'] = AutoPipelineForText2Image.from_pretrained(
                 model['repo_id'],
-                vae=None,
                 torch_dtype=model['compute_type'],
                 safety_checker=None
             ).to(device)