diff --git a/modules_forge/main_entry.py b/modules_forge/main_entry.py index 373c19f9..905ea0f4 100644 --- a/modules_forge/main_entry.py +++ b/modules_forge/main_entry.py @@ -18,7 +18,7 @@ from modules import ( shared_items, ui_common, ) -from modules_forge.presets import PresetArch, use_distill, use_shift +from modules_forge.presets import PresetArch, is_video, use_distill, use_shift logger = logging.getLogger("ui_models") setup_logger(logger) @@ -285,7 +285,10 @@ def on_preset_change(preset: str): else: d_args = {"visible": False} - batch_args = {"minimum": 1, "maximum": 241, "step": 16, "label": "Frames", "value": 1} if preset == "wan" else {"minimum": 1, "maximum": 8, "step": 1, "label": "Batch Size", "value": 1} + if (fps := is_video(preset)) > 1: + batch_args = {"minimum": 1, "maximum": fps * 15 + 1, "step": fps, "label": "Frames", "value": getattr(shared.opts, f"{preset}_batch_size", 1)} + else: + batch_args = {"minimum": 1, "maximum": 8, "step": 1, "label": "Batch Size", "value": getattr(shared.opts, f"{preset}_batch_size", 1)} return [ # ui_checkpoint, ui_vae, ui_forge_unet_dtype diff --git a/modules_forge/presets.py b/modules_forge/presets.py index 2bfdf3b4..681dc562 100644 --- a/modules_forge/presets.py +++ b/modules_forge/presets.py @@ -83,6 +83,10 @@ SHIFT = { PresetArch.ernie: 3.0, } +FRAMES = { + PresetArch.wan.name: 16, +} + def use_distill(arch: str) -> bool: return arch in [preset.name for preset in DISTILL.keys()] @@ -92,6 +96,10 @@ def use_shift(arch: str) -> bool: return arch in [preset.name for preset in SHIFT.keys()] +def is_video(arch: str) -> int: + return FRAMES.get(arch, 1) + + def register(options_templates: dict): from gradio import Dropdown, Slider @@ -189,6 +197,25 @@ def register(options_templates: dict): ) ) + if (fps := FRAMES.get(arch.name, 1)) > 1: + options_templates.update( + options_section( + (f"ui_{name}", name.upper(), "presets"), + { + f"{name}_batch_size": OptionInfo(1, "Frames", Slider, {"minimum": 1, "maximum": fps * 15 + 1, "step": fps}), + }, + ) + ) + else: + options_templates.update( + options_section( + (f"ui_{name}", name.upper(), "presets"), + { + f"{name}_batch_size": OptionInfo(1, "Batch Size", Slider, {"minimum": 1, "maximum": 8, "step": 1}), + }, + ) + ) + options_templates.update( options_section( (f"ui_{name}", name.upper(), "presets"),