refactor: move debiased estim. loss from BaseUI/GeneralUI to Experime…

…ntalArgsUI misc section

main_ui_files/ExperimentalArgsUI.py

@@ -46,6 +46,9 @@ def setup_connections(self) -> None:
         self.widget.cfm_enable.clicked.connect(
             lambda x: self.edit_args("contrastive_flow_matching", x, True)
         )
+        self.widget.debiased_estimation_loss_enable.clicked.connect(
+            lambda x: self.edit_args("debiased_estimation_loss", x, True)
+        )
 
     def edit_args(self, name: str, value: object, optional: bool = False) -> None:
         """Update args dict, handling optional values."""
@@ -80,6 +83,7 @@ def load_args(self, args: dict) -> bool:
 
         # misc args
         self.widget.cfm_enable.setChecked(args.get("contrastive_flow_matching", False))
+        self.widget.debiased_estimation_loss_enable.setChecked(args.get("debiased_estimation_loss", False))
 
 
         # sync args from UI
@@ -104,6 +108,7 @@ def load_args(self, args: dict) -> bool:
 
         # misc args
         self.edit_args("contrastive_flow_matching", self.widget.cfm_enable.isChecked(), True)
+        self.edit_args("debiased_estimation_loss", self.widget.debiased_estimation_loss_enable.isChecked(), True)
 
         return True
 

main_ui_files/GeneralUI.py

@@ -124,9 +124,6 @@ def setup_connections(self) -> None:
         self.widget.v_pred_enable.clicked.connect(
             lambda x: self.edit_args("scale_v_pred_loss_like_noise_pred", x, True)
         )
-        self.widget.debiased_estimation_loss_enable.clicked.connect(
-            lambda x: self.edit_args("debiased_estimation_loss", x, True)
-        )
         self.widget.FP16_enable.clicked.connect(lambda x: self.change_full_type(x, False))
         self.widget.BF16_enable.clicked.connect(lambda x: self.change_full_type(False, x))
         self.widget.FP8_enable.clicked.connect(lambda x: self.edit_args("fp8_base", x, True))
@@ -340,7 +337,6 @@ def load_args(self, args: dict) -> bool:
         self.widget.high_vram_enable.setChecked(args.get("highvram", False))
         self.widget.v_param_enable.setChecked(args.get("v_parameterization", False))
         self.widget.v_pred_enable.setChecked(args.get("scale_v_pred_loss_like_noise_pred", False))
-        self.widget.debiased_estimation_loss_enable.setChecked(args.get("debiased_estimation_loss", False))
         self.widget.FP16_enable.setChecked(args.get("full_fp16", False))
         self.widget.BF16_enable.setChecked(args.get("full_bf16", False))
         self.widget.FP8_enable.setChecked(args.get("fp8_base", False))
@@ -382,8 +378,6 @@ def load_args(self, args: dict) -> bool:
         self.edit_args("lowram", self.widget.low_ram_enable.isChecked(), True)
         self.edit_args("highvram", self.widget.high_vram_enable.isChecked(), True)
         self.enable_disable_v_param(self.widget.v_param_enable.isChecked())
-        # Update args for the new debiased estimation loss checkbox
-        self.edit_args("debiased_estimation_loss", self.widget.debiased_estimation_loss_enable.isChecked(), True)
         self.change_full_type(self.widget.FP16_enable.isChecked(), self.widget.BF16_enable.isChecked())
         self.edit_args("fp8_base", self.widget.FP8_enable.isChecked(), True)
         self.edit_args(

ui_files/BaseUI.py

@@ -27,7 +27,7 @@ class Ui_base_args_ui(object):
     def setupUi(self, base_args_ui):
         if not base_args_ui.objectName():
             base_args_ui.setObjectName(u"base_args_ui")
-        base_args_ui.resize(675, 573)
+        base_args_ui.resize(675, 563)
         self.gridLayout_3 = QGridLayout(base_args_ui)
         self.gridLayout_3.setObjectName(u"gridLayout_3")
         self.formLayout_5 = QFormLayout()
@@ -376,59 +376,52 @@ def setupUi(self, base_args_ui):
         self.model_checkbox_gridLayout = QGridLayout()
         self.model_checkbox_gridLayout.setObjectName(u"model_checkbox_gridLayout")
         self.model_checkbox_gridLayout.setVerticalSpacing(1)
-        self.low_ram_enable = QCheckBox(self.base_model_box)
-        self.low_ram_enable.setObjectName(u"low_ram_enable")
+        self.v_param_enable = QCheckBox(self.base_model_box)
+        self.v_param_enable.setObjectName(u"v_param_enable")
+        self.v_param_enable.setEnabled(True)
 
-        self.model_checkbox_gridLayout.addWidget(self.low_ram_enable, 0, 3, 1, 1)
+        self.model_checkbox_gridLayout.addWidget(self.v_param_enable, 1, 0, 2, 1)
 
         self.v_pred_enable = QCheckBox(self.base_model_box)
         self.v_pred_enable.setObjectName(u"v_pred_enable")
         self.v_pred_enable.setEnabled(False)
 
         self.model_checkbox_gridLayout.addWidget(self.v_pred_enable, 1, 1, 2, 1)
 
-        self.BF16_enable = QCheckBox(self.base_model_box)
-        self.BF16_enable.setObjectName(u"BF16_enable")
+        self.v2_enable = QCheckBox(self.base_model_box)
+        self.v2_enable.setObjectName(u"v2_enable")
 
-        self.model_checkbox_gridLayout.addWidget(self.BF16_enable, 1, 3, 2, 1)
+        self.model_checkbox_gridLayout.addWidget(self.v2_enable, 0, 0, 1, 1)
 
         self.no_half_vae_enable = QCheckBox(self.base_model_box)
         self.no_half_vae_enable.setObjectName(u"no_half_vae_enable")
 
         self.model_checkbox_gridLayout.addWidget(self.no_half_vae_enable, 0, 2, 1, 1)
 
-        self.debiased_estimation_loss_enable = QCheckBox(self.base_model_box)
-        self.debiased_estimation_loss_enable.setObjectName(u"debiased_estimation_loss_enable")
-        sizePolicy1.setHeightForWidth(self.debiased_estimation_loss_enable.sizePolicy().hasHeightForWidth())
-        self.debiased_estimation_loss_enable.setSizePolicy(sizePolicy1)
-
-        self.model_checkbox_gridLayout.addWidget(self.debiased_estimation_loss_enable, 3, 0, 1, 2)
-
         self.FP16_enable = QCheckBox(self.base_model_box)
         self.FP16_enable.setObjectName(u"FP16_enable")
 
         self.model_checkbox_gridLayout.addWidget(self.FP16_enable, 1, 2, 2, 1)
 
-        self.v2_enable = QCheckBox(self.base_model_box)
-        self.v2_enable.setObjectName(u"v2_enable")
+        self.sdxl_enable = QCheckBox(self.base_model_box)
+        self.sdxl_enable.setObjectName(u"sdxl_enable")
 
-        self.model_checkbox_gridLayout.addWidget(self.v2_enable, 0, 0, 1, 1)
+        self.model_checkbox_gridLayout.addWidget(self.sdxl_enable, 0, 1, 1, 1)
+
+        self.low_ram_enable = QCheckBox(self.base_model_box)
+        self.low_ram_enable.setObjectName(u"low_ram_enable")
+
+        self.model_checkbox_gridLayout.addWidget(self.low_ram_enable, 0, 3, 1, 1)
 
         self.FP8_enable = QCheckBox(self.base_model_box)
         self.FP8_enable.setObjectName(u"FP8_enable")
 
         self.model_checkbox_gridLayout.addWidget(self.FP8_enable, 1, 4, 2, 1)
 
-        self.sdxl_enable = QCheckBox(self.base_model_box)
-        self.sdxl_enable.setObjectName(u"sdxl_enable")
-
-        self.model_checkbox_gridLayout.addWidget(self.sdxl_enable, 0, 1, 1, 1)
-
-        self.v_param_enable = QCheckBox(self.base_model_box)
-        self.v_param_enable.setObjectName(u"v_param_enable")
-        self.v_param_enable.setEnabled(True)
+        self.BF16_enable = QCheckBox(self.base_model_box)
+        self.BF16_enable.setObjectName(u"BF16_enable")
 
-        self.model_checkbox_gridLayout.addWidget(self.v_param_enable, 1, 0, 2, 1)
+        self.model_checkbox_gridLayout.addWidget(self.BF16_enable, 1, 3, 2, 1)
 
         self.high_vram_enable = QCheckBox(self.base_model_box)
         self.high_vram_enable.setObjectName(u"high_vram_enable")
@@ -611,45 +604,41 @@ def retranslateUi(self, base_args_ui):
 #endif // QT_CONFIG(tooltip)
         self.vae_selector.setText("")
 #if QT_CONFIG(tooltip)
-        self.low_ram_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Enable this if the trainer is crashing due to running out of system RAM. Typically, this would only be used when interfacing with Google Colab</p></body></html>", None))
+        self.v_param_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>V param, short for V-Paramatarization or commonly knows as 'v-pred', is a noise schedule that some models use. You can set this to train with this noise schedule versus the EDM version of typical SD1.X and SDXL models</p></body></html>", None))
 #endif // QT_CONFIG(tooltip)
-        self.low_ram_enable.setText(QCoreApplication.translate("base_args_ui", u"Low RAM", None))
+        self.v_param_enable.setText(QCoreApplication.translate("base_args_ui", u"V Param", None))
 #if QT_CONFIG(tooltip)
         self.v_pred_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Scales the loss to be in line with EDM</p></body></html>", None))
 #endif // QT_CONFIG(tooltip)
         self.v_pred_enable.setText(QCoreApplication.translate("base_args_ui", u"Scale V pred loss", None))
 #if QT_CONFIG(tooltip)
-        self.BF16_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Train in full BF16. Not compatable with full FP16 or training precision</p></body></html>", None))
+        self.v2_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Select this if you are using an SD2.X based model</p></body></html>", None))
 #endif // QT_CONFIG(tooltip)
-        self.BF16_enable.setText(QCoreApplication.translate("base_args_ui", u"Full BF16", None))
+        self.v2_enable.setText(QCoreApplication.translate("base_args_ui", u"SD2.X Based", None))
 #if QT_CONFIG(tooltip)
         self.no_half_vae_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>This loads the VAE in FP32 or full precision. Roughly doubles VRAM usage for VAE workloads like latent caching, but is sometimes required on older graphics cards</p></body></html>", None))
 #endif // QT_CONFIG(tooltip)
         self.no_half_vae_enable.setText(QCoreApplication.translate("base_args_ui", u"No Half Vae", None))
-#if QT_CONFIG(tooltip)
-        self.debiased_estimation_loss_enable.setToolTip("")
-#endif // QT_CONFIG(tooltip)
-        self.debiased_estimation_loss_enable.setText(QCoreApplication.translate("base_args_ui", u"Debiased Estimation Loss", None))
 #if QT_CONFIG(tooltip)
         self.FP16_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Allows training on full FP16. Not compatable with full BF16 or training precision</p></body></html>", None))
 #endif // QT_CONFIG(tooltip)
         self.FP16_enable.setText(QCoreApplication.translate("base_args_ui", u"Full FP16", None))
 #if QT_CONFIG(tooltip)
-        self.v2_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Select this if you are using an SD2.X based model</p></body></html>", None))
+        self.sdxl_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Select this if you are using an SDXL based model</p></body></html>", None))
 #endif // QT_CONFIG(tooltip)
-        self.v2_enable.setText(QCoreApplication.translate("base_args_ui", u"SD2.X Based", None))
+        self.sdxl_enable.setText(QCoreApplication.translate("base_args_ui", u"SDXL Based", None))
 #if QT_CONFIG(tooltip)
-        self.FP8_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Loads the base model in FP8, which should reduce VRAM usage by roughly half of FP16. Training Precision must be one of FP16 or BF16</p></body></html>", None))
+        self.low_ram_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Enable this if the trainer is crashing due to running out of system RAM. Typically, this would only be used when interfacing with Google Colab</p></body></html>", None))
 #endif // QT_CONFIG(tooltip)
-        self.FP8_enable.setText(QCoreApplication.translate("base_args_ui", u"FP8 Base", None))
+        self.low_ram_enable.setText(QCoreApplication.translate("base_args_ui", u"Low RAM", None))
 #if QT_CONFIG(tooltip)
-        self.sdxl_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Select this if you are using an SDXL based model</p></body></html>", None))
+        self.FP8_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Loads the base model in FP8, which should reduce VRAM usage by roughly half of FP16. Training Precision must be one of FP16 or BF16</p></body></html>", None))
 #endif // QT_CONFIG(tooltip)
-        self.sdxl_enable.setText(QCoreApplication.translate("base_args_ui", u"SDXL Based", None))
+        self.FP8_enable.setText(QCoreApplication.translate("base_args_ui", u"FP8 Base", None))
 #if QT_CONFIG(tooltip)
-        self.v_param_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>V param, short for V-Paramatarization or commonly knows as 'v-pred', is a noise schedule that some models use. You can set this to train with this noise schedule versus the EDM version of typical SD1.X and SDXL models</p></body></html>", None))
+        self.BF16_enable.setToolTip(QCoreApplication.translate("base_args_ui", u"<html><head/><body><p>Train in full BF16. Not compatable with full FP16 or training precision</p></body></html>", None))
 #endif // QT_CONFIG(tooltip)
-        self.v_param_enable.setText(QCoreApplication.translate("base_args_ui", u"V Param", None))
+        self.BF16_enable.setText(QCoreApplication.translate("base_args_ui", u"Full BF16", None))
         self.high_vram_enable.setText(QCoreApplication.translate("base_args_ui", u"High VRAM", None))
     # retranslateUi
 

ui_files/BaseUI.ui

@@ -7,7 +7,7 @@
     <x>0</x>
     <y>0</y>
     <width>675</width>
-    <height>573</height>
+    <height>563</height>
    </rect>
   </property>
   <property name="windowTitle">
@@ -662,13 +662,16 @@
         <property name="verticalSpacing">
          <number>1</number>
         </property>
-        <item row="0" column="3">
-         <widget class="QCheckBox" name="low_ram_enable">
+        <item row="1" column="0" rowspan="2">
+         <widget class="QCheckBox" name="v_param_enable">
+          <property name="enabled">
+           <bool>true</bool>
+          </property>
           <property name="toolTip">
-           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Enable this if the trainer is crashing due to running out of system RAM. Typically, this would only be used when interfacing with Google Colab&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
+           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;V param, short for V-Paramatarization or commonly knows as 'v-pred', is a noise schedule that some models use. You can set this to train with this noise schedule versus the EDM version of typical SD1.X and SDXL models&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
           </property>
           <property name="text">
-           <string>Low RAM</string>
+           <string>V Param</string>
           </property>
          </widget>
         </item>
@@ -685,13 +688,13 @@
           </property>
          </widget>
         </item>
-        <item row="1" column="3" rowspan="2">
-         <widget class="QCheckBox" name="BF16_enable">
+        <item row="0" column="0">
+         <widget class="QCheckBox" name="v2_enable">
           <property name="toolTip">
-           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Train in full BF16. Not compatable with full FP16 or training precision&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
+           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Select this if you are using an SD2.X based model&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
           </property>
           <property name="text">
-           <string>Full BF16</string>
+           <string>SD2.X Based</string>
           </property>
          </widget>
         </item>
@@ -705,22 +708,6 @@
           </property>
          </widget>
         </item>
-        <item row="3" column="0" colspan="2">
-         <widget class="QCheckBox" name="debiased_estimation_loss_enable">
-          <property name="sizePolicy">
-           <sizepolicy hsizetype="Preferred" vsizetype="Fixed">
-            <horstretch>0</horstretch>
-            <verstretch>0</verstretch>
-           </sizepolicy>
-          </property>
-          <property name="toolTip">
-           <string/>
-          </property>
-          <property name="text">
-           <string>Debiased Estimation Loss</string>
-          </property>
-         </widget>
-        </item>
         <item row="1" column="2" rowspan="2">
          <widget class="QCheckBox" name="FP16_enable">
           <property name="toolTip">
@@ -731,46 +718,43 @@
           </property>
          </widget>
         </item>
-        <item row="0" column="0">
-         <widget class="QCheckBox" name="v2_enable">
+        <item row="0" column="1">
+         <widget class="QCheckBox" name="sdxl_enable">
           <property name="toolTip">
-           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Select this if you are using an SD2.X based model&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
+           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Select this if you are using an SDXL based model&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
           </property>
           <property name="text">
-           <string>SD2.X Based</string>
+           <string>SDXL Based</string>
           </property>
          </widget>
         </item>
-        <item row="1" column="4" rowspan="2">
-         <widget class="QCheckBox" name="FP8_enable">
+        <item row="0" column="3">
+         <widget class="QCheckBox" name="low_ram_enable">
           <property name="toolTip">
-           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Loads the base model in FP8, which should reduce VRAM usage by roughly half of FP16. Training Precision must be one of FP16 or BF16&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
+           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Enable this if the trainer is crashing due to running out of system RAM. Typically, this would only be used when interfacing with Google Colab&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
           </property>
           <property name="text">
-           <string>FP8 Base</string>
+           <string>Low RAM</string>
           </property>
          </widget>
         </item>
-        <item row="0" column="1">
-         <widget class="QCheckBox" name="sdxl_enable">
+        <item row="1" column="4" rowspan="2">
+         <widget class="QCheckBox" name="FP8_enable">
           <property name="toolTip">
-           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Select this if you are using an SDXL based model&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
+           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Loads the base model in FP8, which should reduce VRAM usage by roughly half of FP16. Training Precision must be one of FP16 or BF16&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
           </property>
           <property name="text">
-           <string>SDXL Based</string>
+           <string>FP8 Base</string>
           </property>
          </widget>
         </item>
-        <item row="1" column="0" rowspan="2">
-         <widget class="QCheckBox" name="v_param_enable">
-          <property name="enabled">
-           <bool>true</bool>
-          </property>
+        <item row="1" column="3" rowspan="2">
+         <widget class="QCheckBox" name="BF16_enable">
           <property name="toolTip">
-           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;V param, short for V-Paramatarization or commonly knows as 'v-pred', is a noise schedule that some models use. You can set this to train with this noise schedule versus the EDM version of typical SD1.X and SDXL models&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
+           <string>&lt;html&gt;&lt;head/&gt;&lt;body&gt;&lt;p&gt;Train in full BF16. Not compatable with full FP16 or training precision&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>
           </property>
           <property name="text">
-           <string>V Param</string>
+           <string>Full BF16</string>
           </property>
          </widget>
         </item>