Merge pull request #4083 from vladmandic/dev

merge dev

CHANGELOG.md

@@ -1,5 +1,21 @@
 # Change Log for SD.Next
 
+## Update for 2025-07-30
+
+- **Feature**  
+  - Wan select which stage to run: *first/second/both* with configurable *boundary ration* when running both stages  
+    in settings -> model options  
+  - prompt parser allow explict `BOS` and `EOS` tokens in prompt  
+- **UI**  
+  - modernui checkbox/radio styling  
+- **Fixes**  
+  - fix Wan 2.2-5B I2V workflow  
+  - fix inpaint image metadata  
+  - fix processing image save loop  
+  - fix progress bar with refine/detailer  
+  - fix api progress reporting endpoint  
+  - add missing interrogate in output panel  
+
 ## Update for 2025-07-29
 
 ### Highlights for 2025-07-29

installer.py

@@ -593,7 +593,7 @@ def check_diffusers():
     t_start = time.time()
     if args.skip_all or args.skip_git:
         return
-    sha = '56d438727036b0918b30bbe3110c5fe1634ed19d' # diffusers commit hash
+    sha = 'c052791b5fe29ce8a308bf63dda97aa205b729be' # diffusers commit hash
     pkg = pkg_resources.working_set.by_key.get('diffusers', None)
     minor = int(pkg.version.split('.')[1] if pkg is not None else -1)
     cur = opts.get('diffusers_version', '') if minor > -1 else ''

modules/api/server.py

@@ -85,7 +85,7 @@ def get_history(req: models.ReqHistory = Depends()):
     return res
 
 def get_progress(req: models.ReqProgress = Depends()):
-    if shared.state.job_count == 0:
+    if shared.state.job_count == 0: # idle state
         return models.ResProgress(id=shared.state.id, progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo)
     shared.state.do_set_current_image()
     current_image = None
@@ -94,12 +94,17 @@ def get_progress(req: models.ReqProgress = Depends()):
     batch_x = max(shared.state.job_no, 0)
     batch_y = max(shared.state.job_count, 1)
     step_x = max(shared.state.sampling_step, 0)
+    prev_steps = max(shared.state.sampling_steps, 1)
+    while step_x > shared.state.sampling_steps:
+        shared.state.sampling_steps += prev_steps
     step_y = max(shared.state.sampling_steps, 1)
     current = step_y * batch_x + step_x
     total = step_y * batch_y
     progress = min((current / total) if current > 0 and total > 0 else 0, 1)
     time_since_start = time.time() - shared.state.time_start
     eta_relative = (time_since_start / progress) - time_since_start if progress > 0 else 0
+    # shared.log.critical(f'get_progress: batch {batch_x}/{batch_y} step {step_x}/{step_y} current {current}/{total} time={time_since_start} eta={eta_relative}')
+    # shared.log.critical(shared.state)
     res = models.ResProgress(id=shared.state.id, progress=round(progress, 2), eta_relative=round(eta_relative, 2), current_image=current_image, textinfo=shared.state.textinfo, state=shared.state.dict(), )
     return res
 

modules/generation_parameters_copypaste.py

@@ -237,8 +237,10 @@ def connect_paste(button, local_paste_fields, input_comp, override_settings_comp
                     if hasattr(output, "step") and type(output.step) == float:
                         valtype = float
                     debug(f'Paste: "{key}"="{v}" type={valtype} var={vars(output)}')
-                    if valtype == bool and v == "False":
-                        val = False
+                    if valtype == bool:
+                        val = False if v.lower() == "false" else True
+                    elif valtype == list:
+                        val = v if isinstance(v, list) else [item.strip() for item in v.split(',')]
                     else:
                         val = valtype(v)
                     res.append(gr.update(value=val))

modules/img2img.py

@@ -311,9 +311,9 @@ def img2img(id_task: str, state: str, mode: int,
     if mask:
         p.extra_generation_params["Mask blur"] = mask_blur
         p.extra_generation_params["Mask alpha"] = mask_alpha
-        p.extra_generation_params["Mask invert"] = inpainting_mask_invert
-        p.extra_generation_params["Mask area"] = inpaint_full_res
         p.extra_generation_params["Mask padding"] = inpaint_full_res_padding
+        p.extra_generation_params["Mask invert"] = ['masked', 'invert'][inpainting_mask_invert]
+        p.extra_generation_params["Mask area"] = ["full", "masked"][inpaint_full_res]
     p.is_batch = mode == 5
     if p.is_batch:
         process_batch(p, img2img_batch_files, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir, args)

modules/interrogate/interrogate.py

@@ -9,6 +9,7 @@ def interrogate(image):
     if isinstance(image, dict) and 'name' in image:
         image = Image.open(image['name'])
     if image is None:
+        shared.log.error('Interrogate: no image provided')
         return ''
     t0 = time.time()
     if shared.opts.interrogate_default_type == 'OpenCLiP':

modules/model_quant.py

@@ -278,16 +278,16 @@ def load_quanto(msg='', silent=False):
 
 
 def upcast_non_layerwise_modules(model, dtype): # pylint: disable=unused-argument
-    from diffusers.hooks.layerwise_casting import SUPPORTED_PYTORCH_LAYERS
+    from diffusers.hooks.layerwise_casting import _GO_LC_SUPPORTED_PYTORCH_LAYERS
     model_children = list(model.children())
     if not model_children:
-        if not isinstance(model, SUPPORTED_PYTORCH_LAYERS):
+        if not isinstance(model, _GO_LC_SUPPORTED_PYTORCH_LAYERS):
             model = model.to(dtype)
         return model
     for module in model_children:
         has_children = list(module.children())
         if not has_children:
-            if not isinstance(module, SUPPORTED_PYTORCH_LAYERS):
+            if not isinstance(module, _GO_LC_SUPPORTED_PYTORCH_LAYERS):
                 module = module.to(dtype)
         else:
             module = upcast_non_layerwise_modules(module, dtype)

modules/processing.py

@@ -373,37 +373,37 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     image.info["parameters"] = info
                     output_images.append(image)
 
+            for i, image in enumerate(output_images):
                 is_grid = len(output_images) == p.batch_size * p.n_iter + 1 and i == 0
-                for image in output_images:
-                    # resize after
-                    if p.selected_scale_tab_after == 1:
-                        p.width_after, p.height_after = int(image.width * p.scale_by_after), int(image.height * p.scale_by_after)
-                    if p.resize_mode_after != 0 and p.resize_name_after != 'None' and not is_grid:
-                        image = images.resize_image(p.resize_mode_after, image, p.width_after, p.height_after, p.resize_name_after, context=p.resize_context_after)
+                # resize after
+                if p.selected_scale_tab_after == 1:
+                    p.width_after, p.height_after = int(image.width * p.scale_by_after), int(image.height * p.scale_by_after)
+                if p.resize_mode_after != 0 and p.resize_name_after != 'None' and not is_grid:
+                    image = images.resize_image(p.resize_mode_after, image, p.width_after, p.height_after, p.resize_name_after, context=p.resize_context_after)
 
-                    # save images
-                    if shared.opts.samples_save and not p.do_not_save_samples and p.outpath_samples is not None:
-                        info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i)
-                        if isinstance(image, list):
-                            for img in image:
-                                images.save_image(img, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p) # main save image
-                        else:
-                            images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p) # main save image
+                # save images
+                if shared.opts.samples_save and not p.do_not_save_samples and p.outpath_samples is not None:
+                    info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i)
+                    if isinstance(image, list):
+                        for img in image:
+                            images.save_image(img, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p) # main save image
+                    else:
+                        images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p) # main save image
 
-                    if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([shared.opts.save_mask, shared.opts.save_mask_composite, shared.opts.return_mask, shared.opts.return_mask_composite]):
-                        image_mask = p.mask_for_overlay.convert('RGB')
-                        image1 = image.convert('RGBA').convert('RGBa')
-                        image2 = Image.new('RGBa', image.size)
-                        mask = images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L')
-                        image_mask_composite = Image.composite(image1, image2, mask).convert('RGBA')
-                        if shared.opts.save_mask:
-                            images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask")
-                        if shared.opts.save_mask_composite:
-                            images.save_image(image_mask_composite, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask-composite")
-                        if shared.opts.return_mask:
-                            output_images.append(image_mask)
-                        if shared.opts.return_mask_composite:
-                            output_images.append(image_mask_composite)
+                if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([shared.opts.save_mask, shared.opts.save_mask_composite, shared.opts.return_mask, shared.opts.return_mask_composite]):
+                    image_mask = p.mask_for_overlay.convert('RGB')
+                    image1 = image.convert('RGBA').convert('RGBa')
+                    image2 = Image.new('RGBa', image.size)
+                    mask = images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L')
+                    image_mask_composite = Image.composite(image1, image2, mask).convert('RGBA')
+                    if shared.opts.save_mask:
+                        images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask")
+                    if shared.opts.save_mask_composite:
+                        images.save_image(image_mask_composite, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask-composite")
+                    if shared.opts.return_mask:
+                        output_images.append(image_mask)
+                    if shared.opts.return_mask_composite:
+                        output_images.append(image_mask_composite)
 
             timer.process.record('post')
             del samples

modules/processing_diffusers.py

@@ -54,7 +54,7 @@ def restore_state(p: processing.StableDiffusionProcessing):
 
 def process_pre(p: processing.StableDiffusionProcessing):
     from modules import ipadapter, sd_hijack_freeu, para_attention, teacache, hidiffusion, ras, pag, cfgzero, transformer_cache, token_merge, linfusion
-    shared.log.info('Processing apply modifiers')
+    shared.log.info('Processing modifiers: apply')
 
     try:
         # apply-with-unapply
@@ -86,7 +86,7 @@ def process_pre(p: processing.StableDiffusionProcessing):
 
 def process_post(p: processing.StableDiffusionProcessing):
     from modules import ipadapter, hidiffusion, ras, pag, cfgzero, token_merge, linfusion
-    shared.log.info('Processing unapply modifiers')
+    shared.log.info('Processing modifiers: unapply')
 
     try:
         sd_models_compile.check_deepcache(enable=False)

modules/processing_helpers.py

@@ -507,7 +507,7 @@ def update_sampler(p, sd_model, second_pass=False):
 def get_job_name(p, model):
     if hasattr(model, 'pipe'):
         model = model.pipe
-    if hasattr(p, 'xyz'):
+    if getattr(p, 'xyz', False):
         return 'Ignore' # xyz grid handles its own jobs
     if sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE:
         return 'Text'

modules/processing_info.py

@@ -117,29 +117,28 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No
         args["Init image size"] = f"{getattr(p, 'init_img_width', 0)}x{getattr(p, 'init_img_height', 0)}"
         args["Init image hash"] = getattr(p, 'init_img_hash', None)
         args['Image CFG scale'] = p.image_cfg_scale
-        args['Resize scale'] = getattr(p, 'scale_by', None)
         args["Mask weight"] = getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None
         args["Denoising strength"] = getattr(p, 'denoising_strength', None)
+        if args["Size"] != args["Init image size"]:
+            args['Resize scale'] = float(getattr(p, 'scale_by', None)) if getattr(p, 'scale_by', None) != 1 else None
+            args['Resize mode'] = shared.resize_modes[p.resize_mode] if shared.resize_modes[p.resize_mode] != 'None' else None
         if args["Size"] is None:
             args["Size"] = args["Init image size"]
-        # lookup by index
-        if getattr(p, 'resize_mode', None) is not None:
-            args['Resize mode'] = shared.resize_modes[p.resize_mode] if shared.resize_modes[p.resize_mode] != 'None' else None
     if p.resize_mode_before != 0 and p.resize_name_before != 'None' and hasattr(p, 'init_images') and p.init_images is not None and len(p.init_images) > 0:
         args['Resize before'] = f"{p.width_before}x{p.height_before}"
         args['Resize mode before'] = p.resize_mode_before
         args['Resize name before'] = p.resize_name_before
-        args['Resize scale before'] = p.scale_by_before if p.scale_by_before != 1.0 else None
+        args['Resize scale before'] = float(p.scale_by_before) if p.scale_by_before != 1.0 else None
     if p.resize_mode_after != 0 and p.resize_name_after != 'None':
         args['Resize after'] = f"{p.width_after}x{p.height_after}"
         args['Resize mode after'] = p.resize_mode_after
         args['Resize name after'] = p.resize_name_after
-        args['Resize scale after'] = p.scale_by_after if p.scale_by_after != 1.0 else None
+        args['Resize scale after'] = float(p.scale_by_after) if p.scale_by_after != 1.0 else None
     if p.resize_name_mask != 'None' and p.scale_by_mask != 1.0:
         args['Resize mask'] = f"{p.width_mask}x{p.height_mask}"
         args['Resize mode mask'] = p.resize_mode_mask
         args['Resize name mask'] = p.resize_name_mask
-        args['Resize scale mask'] = p.scale_by_mask
+        args['Resize scale mask'] = float(p.scale_by_mask)
     if 'detailer' in p.ops:
         args["Detailer"] = ', '.join(shared.opts.detailer_models) if len(shared.opts.detailer_args) == 0 else shared.opts.detailer_args
         args["Detailer steps"] = p.detailer_steps

modules/prompt_parser_diffusers.py

@@ -360,6 +360,8 @@ def get_prompt_schedule(prompt, steps):
 def get_tokens(pipe, msg, prompt):
     global token_dict, token_type # pylint: disable=global-statement
     if shared.sd_loaded and hasattr(pipe, 'tokenizer') and pipe.tokenizer is not None:
+        prompt = prompt.replace(' BOS ', ' !!!!!!!! ').replace(' EOS ', ' !!!!!!! ')
+        debug(f'Prompt tokenizer: type={msg} prompt="{prompt}"')
         if token_dict is None or token_type != shared.sd_model_type:
             token_type = shared.sd_model_type
             fn = pipe.tokenizer.name_or_path
@@ -375,6 +377,12 @@ def get_tokens(pipe, msg, prompt):
         has_eos_token = pipe.tokenizer.eos_token_id is not None
         ids = pipe.tokenizer(prompt)
         ids = getattr(ids, 'input_ids', [])
+        if has_bos_token and has_eos_token:
+            for i in range(len(ids)):
+                if ids[i] == 21622:
+                    ids[i] = pipe.tokenizer.bos_token_id
+                elif ids[i] == 15203:
+                    ids[i] = pipe.tokenizer.eos_token_id
         tokens = []
         for i in ids:
             try:
@@ -383,7 +391,7 @@ def get_tokens(pipe, msg, prompt):
             except Exception:
                 tokens.append(f'UNK_{i}')
         token_count = len(ids) - int(has_bos_token) - int(has_eos_token)
-        debug(f'Prompt tokenizer: type={msg} tokens={token_count} {tokens}')
+        debug(f'Prompt tokenizer: type={msg} tokens={token_count} tokens={tokens} ids={ids}')
     return token_count
 
 

modules/shared.py

@@ -200,7 +200,8 @@ options_templates.update(options_section(('model_options', "Models Options"), {
     "model_h1_sep": OptionInfo("<h2>HiDream</h2>", "", gr.HTML),
     "model_h1_llama_repo": OptionInfo("Default", "LLama repo", gr.Textbox),
     "model_wan_sep": OptionInfo("<h2>WanAI</h2>", "", gr.HTML),
-    "model_wan_disable_t2": OptionInfo(True, "Disable second stage"),
+    "model_wan_stage": OptionInfo("first", "Processing stage", gr.Radio, {"choices": ['first', 'second', 'both'] }),
+    "model_wan_boundary": OptionInfo(0.85, "Stage boundary ratio", gr.Slider, {"minimum": 0, "maximum": 1.0, "step": 0.05 }),
 }))
 
 options_templates.update(options_section(('vae_encoder', "Variational Auto Encoder"), {

modules/shared_state.py

@@ -178,7 +178,7 @@ class State:
         self.results = []
         self.id = self.get_id(task_id)
         self.job = title
-        self.job_count = 0
+        self.job_count = 1 # cannot be less than 1 on new job
         self.frame_count = 0
         self.batch_no = 0
         self.batch_count = 0
@@ -231,7 +231,7 @@ class State:
             self.sampling_steps = steps
             self.job_count = jobs
         else:
-            self.sampling_steps += steps * jobs
+            self.sampling_steps += (steps * jobs)
             self.job_count += jobs
         self.job = job
         self.history('update')

modules/styles.py

@@ -368,6 +368,9 @@ class StyleDatabase:
             return
         for style in p.styles:
             s = self.find_style(style)
+            if s == self.no_style:
+                shared.log.warning(f'Apply style: name="{style}" not found')
+                continue
             apply_styles_to_extra(p, s)
 
     def extract_comments(self, p):

modules/ui_common.py

@@ -246,7 +246,7 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None, transfe
                                         elem_classes=["gallery_main"],
                                        )
             if prompt is not None:
-                ui_sections.create_interrogate_button(tab=tabname, inputs=result_gallery, outputs=prompt)
+                ui_sections.create_interrogate_button(tab=tabname, inputs=result_gallery, outputs=prompt, what='output')
 
         with gr.Column(elem_id=f"{tabname}_footer", elem_classes="gallery_footer"):
             dummy_component = gr.Label(visible=False)

modules/ui_control.py

@@ -191,7 +191,7 @@ def create_ui(_blocks: gr.Blocks=None):
                             input_image = gr.Image(label="Input", show_label=False, type="pil", interactive=True, tool="editor", height=gr_height, visible=True, image_mode='RGB', elem_id='control_input_select', elem_classes=['control-image'])
                             input_resize = gr.Image(label="Input", show_label=False, type="pil", interactive=True, tool="select", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_resize', elem_classes=['control-image'])
                             input_inpaint = gr.Image(label="Input", show_label=False, type="pil", interactive=True, tool="sketch", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_inpaint', brush_radius=32, mask_opacity=0.6, elem_classes=['control-image'])
-                            btn_interrogate = ui_sections.create_interrogate_button('control')
+                            btn_interrogate = ui_sections.create_interrogate_button('control', what='input')
                             with gr.Row():
                                 input_buttons = [gr.Button('Select', visible=True, interactive=False), gr.Button('Inpaint', visible=True, interactive=True), gr.Button('Outpaint', visible=True, interactive=True)]
                         with gr.Tab('Video', id='in-video') as tab_video:
@@ -619,6 +619,7 @@ def create_ui(_blocks: gr.Blocks=None):
                 # prompt
                 (prompt, "Prompt"),
                 (negative, "Negative prompt"),
+                (styles, "Styles"),
                 # input
                 (denoising_strength, "Denoising strength"),
                 # size basic

modules/ui_img2img.py

@@ -69,7 +69,7 @@ def create_ui():
                     state = gr.Textbox(value='', visible=False)
                     with gr.TabItem('Image', id='img2img_image', elem_id="img2img_image_tab") as tab_img2img:
                         img_init = gr.Image(label="", elem_id="img2img_image", show_label=False, interactive=True, type="pil", tool="editor", image_mode="RGBA", height=512)
-                        interrogate_btn = ui_sections.create_interrogate_button(tab='img2img')
+                        interrogate_btn = ui_sections.create_interrogate_button(tab='img2img', what='input')
                         add_copy_image_controls('img2img', img_init)
 
                     with gr.TabItem('Inpaint', id='img2img_inpaint', elem_id="img2img_inpaint_tab") as tab_inpaint:
@@ -144,7 +144,7 @@ def create_ui():
                             mask_alpha = gr.Slider(label="Alpha", minimum=0.0, maximum=1.0, step=0.05, value=1.0, elem_id="img2img_mask_alpha")
                         with gr.Row():
                             inpainting_mask_invert = gr.Radio(label='Inpaint Mode', choices=['masked', 'invert'], value='masked', type="index", elem_id="img2img_mask_mode")
-                            inpaint_full_res = gr.Radio(label="Inpaint area", choices=["full", "masked"], type="index", value="full", elem_id="img2img_inpaint_full_res")
+                            inpaint_full_res = gr.Radio(label="Inpaint area", choices=["full", "masked"], value="full", type="index", elem_id="img2img_inpaint_full_res")
 
                         def select_img2img_tab(tab):
                             return gr.update(visible=tab in [2, 3, 4]), gr.update(visible=tab == 3)
@@ -236,6 +236,7 @@ def create_ui():
                 # prompt
                 (img2img_prompt, "Prompt"),
                 (img2img_negative_prompt, "Negative prompt"),
+                (img2img_prompt_styles, "Styles"),
                 # sampler
                 (sampler_index, "Sampler"),
                 (steps, "Steps"),
@@ -295,9 +296,9 @@ def create_ui():
                 # inpaint
                 (mask_blur, "Mask blur"),
                 (mask_alpha, "Mask alpha"),
+                (inpaint_full_res_padding, "Mask padding"),
                 (inpainting_mask_invert, "Mask invert"),
                 (inpaint_full_res, "Mask area"),
-                (inpaint_full_res_padding, "Masked padding"),
                 # hidden
                 (seed_resize_from_w, "Seed resize from-1"),
                 (seed_resize_from_h, "Seed resize from-2"),

modules/ui_sections.py

@@ -92,8 +92,8 @@ def create_resolution_inputs(tab, default_width=1024, default_height=1024):
     return width, height
 
 
-def create_interrogate_button(tab: str, inputs: list = None, outputs: str = None):
-    button_interrogate = gr.Button(ui_symbols.interrogate, elem_id=f"{tab}_interrogate", elem_classes=['interrogate'])
+def create_interrogate_button(tab: str, inputs: list = None, outputs: str = None, what: str = ''):
+    button_interrogate = gr.Button(ui_symbols.interrogate, elem_id=f"{tab}_interrogate_{what}", elem_classes=['interrogate'])
     if inputs is not None and outputs is not None:
         button_interrogate.click(fn=interrogate.interrogate, inputs=inputs, outputs=[outputs])
     return button_interrogate

modules/ui_txt2img.py

@@ -92,6 +92,7 @@ def create_ui():
                 # prompt
                 (txt2img_prompt, "Prompt"),
                 (txt2img_negative_prompt, "Negative prompt"),
+                (txt2img_prompt_styles, "Styles"),
                 # main
                 (width, "Size-1"),
                 (height, "Size-2"),

pipelines/model_wanai.py

@@ -8,11 +8,6 @@ def load_transformer(repo_id, diffusers_load_config={}, subfolder='transformer')
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True)
     fn = None
 
-    if subfolder == 'transformer_2' and 'a14b' not in repo_id.lower():
-        return None
-    if subfolder == 'transformer_2' and shared.opts.model_wan_disable_t2:
-        return None
-
     if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
         from modules import sd_unet
         if shared.opts.sd_unet not in list(sd_unet.unet_dict):
@@ -63,13 +58,28 @@ def load_wan(checkpoint_info, diffusers_load_config={}):
     repo_id = sd_models.path_to_repo(checkpoint_info)
     sd_models.hf_auth_check(checkpoint_info)
 
-    transformer = load_transformer(repo_id, diffusers_load_config, 'transformer')
-    transformer_2 = load_transformer(repo_id, diffusers_load_config, 'transformer_2')
+    if 'a14b' in repo_id.lower():
+        if shared.opts.model_wan_stage == 'first':
+            transformer = load_transformer(repo_id, diffusers_load_config, 'transformer')
+            transformer_2 = None
+        elif shared.opts.model_wan_stage == 'second':
+            transformer = load_transformer(repo_id, diffusers_load_config, 'transformer_2')
+            transformer_2 = None
+        elif shared.opts.model_wan_stage == 'both':
+            transformer = load_transformer(repo_id, diffusers_load_config, 'transformer')
+            transformer_2 = load_transformer(repo_id, diffusers_load_config, 'transformer_2')
+        else:
+            shared.log.error(f'Load model: type=WanAI stage="{shared.opts.model_wan_stage}" unsupported')
+            return None
+    else:
+        transformer = load_transformer(repo_id, diffusers_load_config, 'transformer')
+        transformer_2 = None
+
     text_encoder = load_text_encoder(repo_id, diffusers_load_config)
 
     load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
-    boundary_ratio = 0.8 if transformer_2 is not None else None
-    shared.log.debug(f'Load model: type=WanAI model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args} boundary={boundary_ratio}')
+    boundary_ratio = shared.opts.model_wan_boundary if transformer_2 is not None else None
+    shared.log.debug(f'Load model: type=WanAI model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args} stage={shared.opts.model_wan_stage} boundary={boundary_ratio}')
 
     cls = diffusers.WanPipeline
     pipe = cls.from_pretrained(
@@ -88,6 +98,7 @@ def load_wan(checkpoint_info, diffusers_load_config={}):
 
     del text_encoder
     del transformer
+    del transformer_2
 
     sd_hijack_te.init_hijack(pipe)
     from modules.video_models import video_vae