add some fixes to the UI

finetrainers/args.py

@@ -853,8 +853,9 @@ def _validate_dataset_args(args: BaseArgs):
 
 
 def _validate_validation_args(args: BaseArgs):
-    if args.dp_shards > 1 and args.enable_model_cpu_offload:
-        raise ValueError("Model CPU offload is not supported with FSDP at the moment.")
+    if args.enable_model_cpu_offload:
+        if any(x > 1 for x in [args.pp_degree, args.dp_degree, args.dp_shards, args.cp_degree, args.tp_degree]):
+            raise ValueError("Model CPU offload is not supported on multi-GPU at the moment.")
 
 
 def _display_helper_messages(args: argparse.Namespace):

finetrainers/patches/__init__.py

@@ -17,7 +17,12 @@ def perform_patches_for_training(args: "BaseArgs", parallel_backend: "ParallelBa
         if parallel_backend.tensor_parallel_enabled:
             patch.patch_apply_rotary_emb_for_tp_compatibility()
 
+    if args.model_name == ModelType.WAN and "transformer" in args.layerwise_upcasting_modules:
+        from .models.wan import patch
+
+        patch.patch_time_text_image_embedding_forward()
+
     if args.training_type == TrainingType.LORA and len(args.layerwise_upcasting_modules) > 0:
-        from dependencies.peft import patch
+        from .dependencies.peft import patch
 
         patch.patch_peft_move_adapter_to_device_of_base_layer()

finetrainers/patches/models/ltx_video/patch.py

@@ -16,7 +16,7 @@ def patch_apply_rotary_emb_for_tp_compatibility() -> None:
 
 
 def _perform_ltx_transformer_forward_patch() -> None:
-    LTXVideoTransformer3DModel.forward = _patched_LTXVideoTransformer3Dforward
+    LTXVideoTransformer3DModel.forward = _patched_LTXVideoTransformer3D_forward
 
 
 def _perform_ltx_apply_rotary_emb_tensor_parallel_compatibility_patch() -> None:
@@ -35,7 +35,7 @@ def _perform_ltx_apply_rotary_emb_tensor_parallel_compatibility_patch() -> None:
     diffusers.models.transformers.transformer_ltx.apply_rotary_emb = apply_rotary_emb
 
 
-def _patched_LTXVideoTransformer3Dforward(
+def _patched_LTXVideoTransformer3D_forward(
     self,
     hidden_states: torch.Tensor,
     encoder_hidden_states: torch.Tensor,

finetrainers/patches/models/wan/patch.py

@@ -0,0 +1,33 @@
+from typing import Optional
+
+import diffusers
+import torch
+
+
+def patch_time_text_image_embedding_forward() -> None:
+    _patch_time_text_image_embedding_forward()
+
+
+def _patch_time_text_image_embedding_forward() -> None:
+    diffusers.models.transformers.transformer_wan.WanTimeTextImageEmbedding.forward = (
+        _patched_WanTimeTextImageEmbedding_forward
+    )
+
+
+def _patched_WanTimeTextImageEmbedding_forward(
+    self,
+    timestep: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    encoder_hidden_states_image: Optional[torch.Tensor] = None,
+):
+    # Some code has been removed compared to original implementation in Diffusers
+    # Also, timestep is typed as that of encoder_hidden_states
+    timestep = self.timesteps_proj(timestep).type_as(encoder_hidden_states)
+    temb = self.time_embedder(timestep).type_as(encoder_hidden_states)
+    timestep_proj = self.time_proj(self.act_fn(temb))
+
+    encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+    if encoder_hidden_states_image is not None:
+        encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image)
+
+    return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image

finetrainers/trainer/sft_trainer/trainer.py

@@ -334,6 +334,7 @@ class SFTTrainer:
         parallel_backend = self.state.parallel_backend
         train_state = self.state.train_state
         device = parallel_backend.device
+        dtype = self.args.transformer_dtype
 
         memory_statistics = utils.get_memory_statistics()
         logger.info(f"Memory before training start: {json.dumps(memory_statistics, indent=4)}")
@@ -447,8 +448,8 @@ class SFTTrainer:
 
             logger.debug(f"Starting training step ({train_state.step}/{self.args.train_steps})")
 
-            utils.align_device_and_dtype(latent_model_conditions, device, self.args.transformer_dtype)
-            utils.align_device_and_dtype(condition_model_conditions, device, self.args.transformer_dtype)
+            latent_model_conditions = utils.align_device_and_dtype(latent_model_conditions, device, dtype)
+            condition_model_conditions = utils.align_device_and_dtype(condition_model_conditions, device, dtype)
             latent_model_conditions = utils.make_contiguous(latent_model_conditions)
             condition_model_conditions = utils.make_contiguous(condition_model_conditions)
 
@@ -729,6 +730,7 @@ class SFTTrainer:
 
         parallel_backend.wait_for_everyone()
         if not final_validation:
+            self._move_components_to_device()
             self.transformer.train()
 
     def _evaluate(self) -> None:

vms/ui/app_ui.py

@@ -146,8 +146,8 @@ class AppUI:
                 # Sidebar for navigation
                 with gr.Sidebar(position="left", open=True):
                     gr.Markdown("# 🎞️ Video Model Studio")
-                    self.components["current_project_btn"] = gr.Button("Current Project", variant="primary")
-                    self.components["system_monitoring_btn"] = gr.Button("System Monitoring")
+                    self.components["current_project_btn"] = gr.Button("📂 Current Project", variant="primary")
+                    self.components["system_monitoring_btn"] = gr.Button("🌡️ System Monitoring")
 
                 # Main content area with tabs
                 with gr.Column():
@@ -174,7 +174,7 @@ class AppUI:
                                     tab_obj.create(project_tabs)
                         
                         # Monitoring View Tab
-                        with gr.Tab("📊 System Monitoring", id=1) as monitoring_view:
+                        with gr.Tab("🌡️ System Monitoring", id=1) as monitoring_view:
                             # Create monitoring tabs
                             with gr.Tabs() as monitoring_tabs:
                                 # Store reference to monitoring tabs component
@@ -257,19 +257,13 @@ class AppUI:
             self.project_tabs["train_tab"].components["stop_btn"],
             self.project_tabs["train_tab"].components["delete_checkpoints_btn"]
         ]
-        
+
         button_timer.tick(
             fn=self.project_tabs["train_tab"].get_button_updates,
             outputs=button_outputs
         )
         
-        
-        # Add delete_checkpoints_btn or pause_resume_btn as the third button
-        if "delete_checkpoints_btn" in self.project_tabs["train_tab"].components:
-            button_outputs.append(self.project_tabs["train_tab"].components["delete_checkpoints_btn"])
-        elif "pause_resume_btn" in self.project_tabs["train_tab"].components:
-            button_outputs.append(self.project_tabs["train_tab"].components["pause_resume_btn"])
-        
+    
         # Dataset refresh timer (every 5 seconds)
         dataset_timer = gr.Timer(value=5)
         dataset_timer.tick(
@@ -558,21 +552,21 @@ class AppUI:
             
             if is_training:
                 # Active training detected
-                start_btn_props = {"interactive": False, "variant": "secondary", "value": "Start new training"}
-                resume_btn_props = {"interactive": False, "variant": "secondary", "value": "Start from latest checkpoint"}
+                start_btn_props = {"interactive": False, "variant": "secondary", "value": "🚀 Start new training"}
+                resume_btn_props = {"interactive": False, "variant": "secondary", "value": "🛰️ Start from latest checkpoint"}
                 stop_btn_props = {"interactive": True, "variant": "primary", "value": "Stop at Last Checkpoint"}
                 delete_btn_props = {"interactive": False, "variant": "stop", "value": "Delete All Checkpoints"}
             else:
                 # No active training
-                start_btn_props = {"interactive": True, "variant": "primary", "value": "Start new training"}
-                resume_btn_props = {"interactive": has_checkpoints, "variant": "primary", "value": "Start from latest checkpoint"}
+                start_btn_props = {"interactive": True, "variant": "primary", "value": "🚀 Start new training"}
+                resume_btn_props = {"interactive": has_checkpoints, "variant": "primary", "value": "🛰️ Start from latest checkpoint"}
                 stop_btn_props = {"interactive": False, "variant": "secondary", "value": "Stop at Last Checkpoint"}
                 delete_btn_props = {"interactive": has_checkpoints, "variant": "stop", "value": "Delete All Checkpoints"}
         else:
             # Use button states from recovery, adding the new resume button
-            start_btn_props = ui_updates.get("start_btn", {"interactive": True, "variant": "primary", "value": "Start new training"})
+            start_btn_props = ui_updates.get("start_btn", {"interactive": True, "variant": "primary", "value": "🚀 Start new training"})
             resume_btn_props = {"interactive": has_checkpoints and not self.training.is_training_running(), 
-                            "variant": "primary", "value": "Start from latest checkpoint"}
+                            "variant": "primary", "value": "🛰️ Start from latest checkpoint"}
             stop_btn_props = ui_updates.get("stop_btn", {"interactive": False, "variant": "secondary", "value": "Stop at Last Checkpoint"})
             delete_btn_props = ui_updates.get("delete_checkpoints_btn", {"interactive": has_checkpoints, "variant": "stop", "value": "Delete All Checkpoints"})
         
@@ -604,7 +598,7 @@ class AppUI:
         
         return (
             gr.Markdown(value=caption_title),
-            gr.Markdown(value=f"{train_title} available for training")
+            gr.Markdown(value=f"{train_title}")
         )
     
     def refresh_dataset(self):

vms/ui/monitoring/tabs/general_tab.py

@@ -32,7 +32,7 @@ class GeneralTab(BaseTab):
         """Create the Monitor tab UI components"""
         with gr.TabItem(self.title, id=self.id) as tab:
             with gr.Row():
-                gr.Markdown("## System Monitoring")
+                gr.Markdown("## 🌡️ System Monitoring")
   
             # Current metrics
             with gr.Row():

vms/ui/project/tabs/manage_tab.py

@@ -29,23 +29,23 @@ class ManageTab(BaseTab):
         with gr.TabItem(self.title, id=self.id) as tab:
             with gr.Row():
                 with gr.Column():
-                    gr.Markdown("## Download your model")
+                    gr.Markdown("## 🏦 Backup your model")
                     gr.Markdown("There is currently a bug, you might have to click multiple times to trigger a download.")
 
                     with gr.Row():
                         self.components["download_dataset_btn"] = gr.DownloadButton(
-                            "Download training dataset",
+                            "📦 Download training dataset (.zip)",
                             variant="secondary",
                             size="lg"
                         )
                         self.components["download_model_btn"] = gr.DownloadButton(
-                            "Download model weights",
+                            "🧠 Download weights (.safetensors)",
                             variant="secondary",
                             size="lg"
                         )
             with gr.Row():
                 with gr.Column():
-                    gr.Markdown("## Publish your model")
+                    gr.Markdown("## 📡 Publish your model")
                     gr.Markdown("You model can be pushed to Hugging Face (this will use HF_API_TOKEN)")
 
             with gr.Row():
@@ -65,19 +65,19 @@ class ManageTab(BaseTab):
 
             with gr.Row():
                 with gr.Column():
-                    gr.Markdown("## Delete your data")
+                    gr.Markdown("## ♻️ Delete your data")
                     gr.Markdown("Make sure you have made a backup first.")
                     gr.Markdown("If you are deleting because of a bug, remember you can use the Developer Mode on HF to inspect the working directory (in /data or .data)")
 
             with gr.Row():
                 with gr.Column():
-                    gr.Markdown("### Delete specific data")
+                    gr.Markdown("### 🧽 Delete specific data")
                     gr.Markdown("You can selectively delete either the dataset and/or the last model data.")
 
             with gr.Row():
                 with gr.Column(scale=1):
                     self.components["delete_dataset_btn"] = gr.Button(
-                        "Delete dataset (images, video, captions)",
+                        "🚨 Delete dataset (images, video, captions)",
                         variant="secondary"
                     )
                     self.components["delete_dataset_status"] = gr.Textbox(
@@ -88,7 +88,7 @@ class ManageTab(BaseTab):
                 
                 with gr.Column(scale=1):
                     self.components["delete_model_btn"] = gr.Button(
-                        "Delete model (checkpoints, weights, config)",
+                        "🚨 Delete model (checkpoints, weights, config)",
                         variant="secondary"
                     )
                     self.components["delete_model_status"] = gr.Textbox(
@@ -99,12 +99,12 @@ class ManageTab(BaseTab):
                 
             with gr.Row():
                 with gr.Column():
-                    gr.Markdown("### Delete everything")
-                    gr.Markdown("This will delete both the dataset (all images, videos and captions) AND the latest model (weights, checkpoints, settings). So use with care!")
+                    gr.Markdown("### ☢️ Nuke all project data")
+                    gr.Markdown("This will nuke the original dataset (all images, videos and captions), the training dataset, and the model outputs (weights, checkpoints, settings). So use with care!")
 
             with gr.Row():
                 self.components["global_stop_btn"] = gr.Button(
-                    "Stop everything and delete my data",
+                    "🚨 Delete all project data and models (are you sure?!)",
                     variant="stop"
                 )
                 self.components["global_status"] = gr.Textbox(

vms/ui/project/tabs/preview_tab.py

@@ -29,7 +29,7 @@ class PreviewTab(BaseTab):
         """Create the Preview tab UI components"""
         with gr.TabItem(self.title, id=self.id) as tab:
             with gr.Row():
-                gr.Markdown("## Preview your model")
+                gr.Markdown("## 🔬 Preview your model")
             
             with gr.Row():
                 with gr.Column(scale=2):
@@ -202,11 +202,11 @@ class PreviewTab(BaseTab):
                         interactive=False
                     )
                     
-                    with gr.Accordion("Log", open=True):
+                    with gr.Accordion("Log", open=False):
                         self.components["log"] = gr.TextArea(
                             label="Generation Log",
                             interactive=False,
-                            lines=20
+                            lines=60
                         )
         
         return tab

vms/ui/project/tabs/train_tab.py

@@ -44,7 +44,7 @@ class TrainTab(BaseTab):
             with gr.Row():
                 with gr.Column():
                     with gr.Row():
-                        self.components["train_title"] = gr.Markdown("## 0 files available for training (0 bytes)")
+                        self.components["train_title"] = gr.Markdown("## 0 files in the training dataset")
 
                     with gr.Row():
                         with gr.Column():
@@ -181,79 +181,80 @@ class TrainTab(BaseTab):
 
                 with gr.Row():
                     with gr.Column():
-                        # Add description of the training buttons
-                        self.components["training_buttons_info"] = gr.Markdown("""
-                        ## Training Options
-                        - **Start new training**: Begins training from scratch (clears previous checkpoints)
-                        - **Start from latest checkpoint**: Continues training from the most recent checkpoint
-                        """)
-                        
-                        with gr.Row():
-                            # Check for existing checkpoints to determine button text
-                            checkpoints = list(OUTPUT_PATH.glob("finetrainers_step_*"))
-                            has_checkpoints = len(checkpoints) > 0
 
-                            # Rename "Start Training" to "Start new training"
-                            self.components["start_btn"] = gr.Button(
-                                "Start new training",
-                                variant="primary",
-                                interactive=not ASK_USER_TO_DUPLICATE_SPACE
-                            )
-                            
-                            # Add new button for continuing from checkpoint
-                            self.components["resume_btn"] = gr.Button(
-                                "Start from latest checkpoint",
-                                variant="primary", 
-                                interactive=has_checkpoints and not ASK_USER_TO_DUPLICATE_SPACE
-                            )
-                            
-                        with gr.Row():
-                            # Just use stop and pause buttons for now to ensure compatibility
-                            self.components["stop_btn"] = gr.Button(
-                                "Stop at Last Checkpoint",
-                                variant="primary",
-                                interactive=False
-                            )
-                            
-                            self.components["pause_resume_btn"] = gr.Button(
-                                "Resume Training",
-                                variant="secondary",
-                                interactive=False,
-                                visible=False
-                            )
-                            
-                            # Add delete checkpoints button
-                            self.components["delete_checkpoints_btn"] = gr.Button(
-                                "Delete All Checkpoints",
-                                variant="stop",
-                                interactive=has_checkpoints
-                            )
+                        with gr.Row():  
+                            with gr.Column():
+                                # Add description of the training buttons
+                                self.components["training_buttons_info"] = gr.Markdown("""
+                                ## ⚗️ Train your model on your dataset
+                                - **Start new training**: Begins training from scratch (clears previous checkpoints)
+                                - **Start from latest checkpoint**: Continues training from the most recent checkpoint
+                                """)
+                                
+                                with gr.Row():
+                                    # Check for existing checkpoints to determine button text
+                                    checkpoints = list(OUTPUT_PATH.glob("finetrainers_step_*"))
+                                    has_checkpoints = len(checkpoints) > 0
 
-                with gr.Column():
-                    with gr.Row():
-                        with gr.Column():
-                            self.components["status_box"] = gr.Textbox(
-                                label="Training Status",
-                                interactive=False,
-                                lines=4
-                            )
-                            
-                            # Add new component for current task progress
-                            self.components["current_task_box"] = gr.Textbox(
-                                label="Current Task Progress",
-                                interactive=False,
-                                lines=3,
-                                elem_id="current_task_display"
-                            )
-                            
-                            with gr.Accordion("Finetrainers output (or see app logs for more details)"):
-                                self.components["log_box"] = gr.TextArea(
-                                    #label="",
+                                    self.components["start_btn"] = gr.Button(
+                                        "🚀 Start new training",
+                                        variant="primary",
+                                        interactive=not ASK_USER_TO_DUPLICATE_SPACE
+                                    )
+                                    
+                                    # Add new button for continuing from checkpoint
+                                    self.components["resume_btn"] = gr.Button(
+                                        "🛰️ Start from latest checkpoint",
+                                        variant="primary", 
+                                        interactive=has_checkpoints and not ASK_USER_TO_DUPLICATE_SPACE
+                                    )
+                                    
+                                with gr.Row():
+                                    # Just use stop and pause buttons for now to ensure compatibility
+                                    self.components["stop_btn"] = gr.Button(
+                                        "Stop at Last Checkpoint",
+                                        variant="primary",
+                                        interactive=False
+                                    )
+                                    
+                                    self.components["pause_resume_btn"] = gr.Button(
+                                        "Resume Training",
+                                        variant="secondary",
+                                        interactive=False,
+                                        visible=False
+                                    )
+                                    
+                                    # Add delete checkpoints button
+                                    self.components["delete_checkpoints_btn"] = gr.Button(
+                                        "Delete All Checkpoints",
+                                        variant="stop",
+                                        interactive=has_checkpoints
+                                    )
+
+                        with gr.Row():
+                            with gr.Column():
+                                self.components["status_box"] = gr.Textbox(
+                                    label="Training Status",
+                                    interactive=False,
+                                    lines=4
+                                )
+                                
+                                # Add new component for current task progress
+                                self.components["current_task_box"] = gr.Textbox(
+                                    label="Current Task Progress",
                                     interactive=False,
-                                    lines=60,
-                                    max_lines=600,
-                                    autoscroll=True
+                                    lines=3,
+                                    elem_id="current_task_display"
                                 )
+                                
+                                with gr.Accordion("Finetrainers output (or see app logs for more details)", open=False):
+                                    self.components["log_box"] = gr.TextArea(
+                                        #label="",
+                                        interactive=False,
+                                        lines=60,
+                                        max_lines=600,
+                                        autoscroll=True
+                                    )
                     
         return tab
     
@@ -963,7 +964,7 @@ class TrainTab(BaseTab):
         
         # Create button updates
         start_btn = gr.Button(
-            value="Start new training",
+            value="🚀 Start new training",
             interactive=not is_training,
             variant="primary" if not is_training else "secondary"
         )