Merge pull request #1841 from bghira/issue/1811

bghira · web-flow · commit a3ecc2d874ad · 2025-10-25T21:05:40.000-06:00
(#1811) disable prompt management options when i2v is selected
diff --git a/simpletuner/helpers/models/common.py b/simpletuner/helpers/models/common.py
@@ -151,6 +151,8 @@ class ModelFoundation(ABC):
     MAXIMUM_CANVAS_SIZE = None
     SUPPORTS_LORA = None
     SUPPORTS_CONTROLNET = None
+    STRICT_I2V_FLAVOURS = tuple()
+    STRICT_I2V_FOR_ALL_FLAVOURS = False
 
     def __init__(self, config: dict, accelerator):
         self.config = config
@@ -181,6 +183,49 @@ def supports_controlnet(cls) -> bool:
             return bool(cls.SUPPORTS_CONTROLNET)
         return False
 
+    @classmethod
+    def strict_i2v_flavours(cls):
+        """
+        Return flavour identifiers that require strict image-to-video validation inputs.
+        """
+        flavours = getattr(cls, "STRICT_I2V_FLAVOURS", tuple())
+        if not flavours:
+            return []
+        if isinstance(flavours, (list, tuple, set, frozenset)):
+            result = []
+            for entry in flavours:
+                try:
+                    value = str(entry).strip()
+                except Exception:
+                    continue
+                if value:
+                    result.append(value)
+            return list(dict.fromkeys(result))
+        if isinstance(flavours, str):
+            trimmed = flavours.strip()
+            return [trimmed] if trimmed else []
+        return []
+
+    @classmethod
+    def is_strict_i2v_flavour(cls, flavour) -> bool:
+        """
+        Determine whether the provided flavour requires strict image-to-video validation inputs.
+        """
+        if getattr(cls, "STRICT_I2V_FOR_ALL_FLAVOURS", False):
+            return True
+        if not flavour:
+            return False
+        try:
+            candidate = str(flavour).strip().lower()
+        except Exception:
+            return False
+        if not candidate:
+            return False
+        for entry in cls.strict_i2v_flavours():
+            if entry.strip().lower() == candidate:
+                return True
+        return False
+
     def log_model_devices(self):
         """
         Log the devices of the model components.
diff --git a/simpletuner/helpers/models/wan/model.py b/simpletuner/helpers/models/wan/model.py
@@ -316,6 +316,7 @@ class Wan(VideoModelFoundation):
             "ti2v-5b-2.2",
         }
     )
+    STRICT_I2V_FLAVOURS = tuple(sorted((I2V_FLAVOURS | FLF2V_FLAVOURS)))
 
     def __init__(self, config, accelerator):
         super().__init__(config, accelerator)
diff --git a/simpletuner/simpletuner_sdk/server/routes/datasets.py b/simpletuner/simpletuner_sdk/server/routes/datasets.py
@@ -113,16 +113,24 @@ async def get_dataset_plan() -> DatasetPlanResponse:
 
     # Get model_family from active config
     model_family = None
+    model_flavour = None
     try:
         from simpletuner.simpletuner_sdk.server.services.configs_service import ConfigsService
 
         configs_service = ConfigsService()
         active_config = configs_service.get_active_config()
-        model_family = active_config["config"].get("model_family") or active_config["config"].get("--model_family")
+        config_blob = active_config["config"]
+        model_family = config_blob.get("model_family") or config_blob.get("--model_family")
+        model_flavour = config_blob.get("model_flavour") or config_blob.get("--model_flavour")
     except Exception:
         pass
 
-    validations = compute_validations(datasets, get_dataset_blueprints(), model_family=model_family)
+    validations = compute_validations(
+        datasets,
+        get_dataset_blueprints(),
+        model_family=model_family,
+        model_flavour=model_flavour,
+    )
     return DatasetPlanResponse(
         datasets=datasets,
         validations=validations,
@@ -149,16 +157,24 @@ def _persist_plan(payload: DatasetPlanPayload) -> DatasetPlanResponse:
 
     # Get model_family from active config
     model_family = None
+    model_flavour = None
     try:
         from simpletuner.simpletuner_sdk.server.services.configs_service import ConfigsService
 
         configs_service = ConfigsService()
         active_config = configs_service.get_active_config()
-        model_family = active_config["config"].get("model_family") or active_config["config"].get("--model_family")
+        config_blob = active_config["config"]
+        model_family = config_blob.get("model_family") or config_blob.get("--model_family")
+        model_flavour = config_blob.get("model_flavour") or config_blob.get("--model_flavour")
     except Exception:
         pass
 
-    validations = compute_validations(datasets, get_dataset_blueprints(), model_family=model_family)
+    validations = compute_validations(
+        datasets,
+        get_dataset_blueprints(),
+        model_family=model_family,
+        model_flavour=model_flavour,
+    )
     errors = [message for message in validations if message.level == "error"]
     if errors:
         raise HTTPException(
@@ -460,17 +476,25 @@ async def create_dataset(dataset: Dict[str, Any]) -> Dict[str, Any]:
 
     # Get model_family from active config
     model_family = None
+    model_flavour = None
     try:
         from simpletuner.simpletuner_sdk.server.services.configs_service import ConfigsService
 
         configs_service = ConfigsService()
         active_config = configs_service.get_active_config()
-        model_family = active_config["config"].get("model_family") or active_config["config"].get("--model_family")
+        config_blob = active_config["config"]
+        model_family = config_blob.get("model_family") or config_blob.get("--model_family")
+        model_flavour = config_blob.get("model_flavour") or config_blob.get("--model_flavour")
     except Exception:
         pass
 
     # Validate the updated plan
-    validations = compute_validations(datasets, get_dataset_blueprints(), model_family=model_family)
+    validations = compute_validations(
+        datasets,
+        get_dataset_blueprints(),
+        model_family=model_family,
+        model_flavour=model_flavour,
+    )
     errors = [v for v in validations if v.level == "error"]
 
     if errors:
@@ -512,17 +536,25 @@ async def update_dataset(dataset_id: str, dataset: Dict[str, Any]) -> Dict[str,
 
     # Get model_family from active config
     model_family = None
+    model_flavour = None
     try:
         from simpletuner.simpletuner_sdk.server.services.configs_service import ConfigsService
 
         configs_service = ConfigsService()
         active_config = configs_service.get_active_config()
-        model_family = active_config.get("model_family") or active_config.get("--model_family")
+        config_blob = active_config["config"]
+        model_family = config_blob.get("model_family") or config_blob.get("--model_family")
+        model_flavour = config_blob.get("model_flavour") or config_blob.get("--model_flavour")
     except Exception:
         pass
 
     # Validate the updated plan
-    validations = compute_validations(datasets, get_dataset_blueprints(), model_family=model_family)
+    validations = compute_validations(
+        datasets,
+        get_dataset_blueprints(),
+        model_family=model_family,
+        model_flavour=model_flavour,
+    )
     errors = [v for v in validations if v.level == "error"]
 
     if errors:
diff --git a/simpletuner/simpletuner_sdk/server/services/dataset_plan.py b/simpletuner/simpletuner_sdk/server/services/dataset_plan.py
@@ -89,6 +89,7 @@ def compute_validations(
     datasets: List[Dict[str, Any]],
     blueprints: Optional[List[BackendBlueprint]] = None,
     model_family: Optional[str] = None,
+    model_flavour: Optional[str] = None,
 ) -> List[ValidationMessage]:
     """Perform lightweight validation mirroring the UI logic."""
     validations: List[ValidationMessage] = []
@@ -149,6 +150,7 @@ def compute_validations(
 
     # Check if this is a video-only model
     is_video_model = False
+    requires_strict_video_inputs = False
     if model_family:
         try:
             from simpletuner.helpers.models.common import VideoModelFoundation
@@ -157,6 +159,8 @@ def compute_validations(
             model_cls = ModelRegistry.get(model_family)
             if model_cls:
                 is_video_model = issubclass(model_cls, VideoModelFoundation)
+                if hasattr(model_cls, "is_strict_i2v_flavour") and callable(model_cls.is_strict_i2v_flavour):
+                    requires_strict_video_inputs = bool(model_cls.is_strict_i2v_flavour(model_flavour))
         except Exception:
             pass
 
@@ -167,13 +171,22 @@ def compute_validations(
 
     if is_video_model:
         if video_count == 0:
-            validations.append(
-                ValidationMessage(
-                    field="datasets",
-                    message="at least one video dataset is required for video models",
-                    level="error",
+            if requires_strict_video_inputs:
+                validations.append(
+                    ValidationMessage(
+                        field="datasets",
+                        message="strict image-to-video flavours require at least one video dataset",
+                        level="error",
+                    )
+                )
+            elif image_count == 0:
+                validations.append(
+                    ValidationMessage(
+                        field="datasets",
+                        message="add at least one image or video dataset for this model",
+                        level="error",
+                    )
                 )
-            )
     else:
         if image_count == 0:
             validations.append(
diff --git a/simpletuner/simpletuner_sdk/server/services/models_service.py b/simpletuner/simpletuner_sdk/server/services/models_service.py
@@ -154,6 +154,14 @@ def get_model_details(self, model_family: str) -> Dict[str, Any]:
             ),
             "is_video_model": issubclass(model_cls, VideoModelFoundation),
         }
+        strict_i2v_flavours: list[str] = []
+        try:
+            if hasattr(model_cls, "strict_i2v_flavours") and callable(model_cls.strict_i2v_flavours):
+                strict_i2v_flavours = list(model_cls.strict_i2v_flavours())
+        except Exception:
+            strict_i2v_flavours = []
+        capabilities["strict_i2v_flavours"] = strict_i2v_flavours
+        capabilities["strict_i2v_all_flavours"] = bool(getattr(model_cls, "STRICT_I2V_FOR_ALL_FLAVOURS", False))
 
         default_flavour = getattr(model_cls, "DEFAULT_MODEL_FLAVOUR", None)
         if default_flavour is None:
diff --git a/simpletuner/static/css/trainer.css b/simpletuner/static/css/trainer.css
@@ -1237,3 +1237,42 @@ optgroup option {
 .config-card:first-child::before {
     display: none;
 }
+
+.prompt-overlay {
+    display: none;
+    position: absolute;
+    inset: 0;
+    background: rgba(12, 15, 25, 0.55);
+    border-radius: 0.75rem;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    padding: 1.25rem;
+    z-index: 2;
+    backdrop-filter: blur(0.4rem);
+}
+
+.form-section.prompt-management-disabled {
+    position: relative;
+}
+
+.form-section.prompt-management-disabled .prompt-overlay {
+    display: flex;
+}
+
+.form-section.prompt-management-disabled .prompt-overlay-content {
+    display: flex;
+    gap: 0.75rem;
+    color: #fff;
+    text-align: left;
+}
+
+.form-section.prompt-management-disabled .prompt-overlay-content strong {
+    display: block;
+    margin-bottom: 0.25rem;
+}
+
+.form-section.prompt-management-disabled .prompt-overlay-content i {
+    font-size: 1.5rem;
+}
+}
diff --git a/simpletuner/templates/components/action_buttons_htmx.html b/simpletuner/templates/components/action_buttons_htmx.html
@@ -97,7 +97,7 @@
                 data-bs-toggle="modal"
                 data-bs-target="#shutdownConfirmModal">
             <i class="fas fa-power-off me-1"></i>
-            <span class="btn-text">&#x23FB; Shutdown</span>
+            <span class="btn-text">Shutdown</span>
         </button>
     </div>
 </div>
diff --git a/simpletuner/templates/form_tab.html b/simpletuner/templates/form_tab.html
diff --git a/simpletuner/templates/trainer_dataloader_section.html b/simpletuner/templates/trainer_dataloader_section.html
diff --git a/simpletuner/templates/trainer_htmx.html b/simpletuner/templates/trainer_htmx.html

Original file line number	Diff line number	Diff line change
`@@ -316,6 +316,7 @@ class Wan(VideoModelFoundation):`
`316`	`316`	`"ti2v-5b-2.2",`
`317`	`317`	`}`
`318`	`318`	`)`
	`319`	`+ STRICT_I2V_FLAVOURS = tuple(sorted((I2V_FLAVOURS \| FLF2V_FLAVOURS)))`
`319`	`320`
`320`	`321`	`def __init__(self, config, accelerator):`
`321`	`322`	`super().__init__(config, accelerator)`