New UI

js/module_dreambooth.js

@@ -238,12 +238,12 @@ function onDbUpdate(data) {
                     x: [stats.session_step],
                     mode: 'lines',
                     name: 'LR',
-                    yaxis: 'y'
+                    yaxis: 'y2'
                 }
             ];
 
             // Specify the layout options for the plot
-            let lr = parseFloat(stats.unet_lr);
+            let lr = parseFloat(stats.lr_data[0]);
             console.log("LR: ", lr);
             let layout = {
                 autosize: true,
@@ -276,7 +276,7 @@ function onDbUpdate(data) {
             let existingData = plotter.data;
 
             // Extend the existing traces by adding new points for loss, VRAM usage, and learning rate
-            let lr = parseFloat(stats.unet_lr);
+            let lr = parseFloat(stats.lr_data[0]);
             console.log("LR: ", lr);
             Plotly.extendTraces(plotter, {
                 y: [[stats.loss], [lr]],

module_dreambooth.py

@@ -29,10 +29,11 @@ logger = logging.getLogger(__name__)
 class DreamboothModule(BaseModule):
 
     def __init__(self):
+        self.id = "dreambooth"
         self.name: str = "Dreambooth"
         self.path = os.path.abspath(os.path.dirname(__file__))
         self.model_handler = ModelHandler()
-        super().__init__(self.name, self.path)
+        super().__init__(self.id, self.name, self.path)
 
     def initialize(self, app: FastAPI, handler: SocketHandler):
         self._initialize_api(app)