Minor tweaks for style create/delete

Acly · Sep 22, 2023 · b0ca44e · b0ca44e
1 parent 6adba37
commit b0ca44e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 17 deletions.
diff --git a/ai_diffusion/style.py b/ai_diffusion/style.py
@@ -97,7 +97,7 @@ class StyleSettings:
 
     cfg_scale = Setting(
         "Guidance Strength (CFG Scale)",
-        7.0,
+        7,
         "Value which indicates how closely image generation follows the text prompt",
     )
 
@@ -193,22 +193,24 @@ def __init__(self, folder: Path = default_folder):
     @property
     def default(self):
         return self[0]
-    
-    def create(self, name: str = "style") -> Style:
+
+    def create(self, name: str = "style", checkpoint: str = "") -> Style:
         if Path(self.folder / f"{name}.json").exists():
             i = 1
             basename = name
             while Path(self.folder / f"{basename}_{i}.json").exists():
                 i += 1
             name = f"{basename}_{i}"
-        
+
         new_style = Style(self.folder / f"{name}.json")
-        new_style.name = name
+        new_style.name = "New Style"
+        if checkpoint:
+            new_style.sd_checkpoint = checkpoint
         self._list.append(new_style)
         new_style.save()
         self.changed.emit()
         return new_style
-    
+
     def delete(self, style: Style):
         self._list.remove(style)
         style.filepath.unlink()

diff --git a/ai_diffusion/ui/settings.py b/ai_diffusion/ui/settings.py
@@ -568,6 +568,16 @@ def __init__(self):
         self._style_list.currentIndexChanged.connect(self._change_style)
         frame_layout.addWidget(self._style_list)
 
+        self._create_style_button = QToolButton(self)
+        self._create_style_button.setIcon(Krita.instance().icon("list-add"))
+        self._create_style_button.clicked.connect(self._create_style)
+        frame_layout.addWidget(self._create_style_button)
+
+        self._delete_style_button = QToolButton(self)
+        self._delete_style_button.setIcon(Krita.instance().icon("deletelayer"))
+        self._delete_style_button.clicked.connect(self._delete_style)
+        frame_layout.addWidget(self._delete_style_button)
+
         self._refresh_button = QToolButton(self)
         self._refresh_button.setIcon(Krita.instance().icon("reload-preset"))
         self._refresh_button.clicked.connect(self._update_style_list)
@@ -577,16 +587,6 @@ def __init__(self):
         self._open_folder_button.setIcon(Krita.instance().icon("document-open"))
         self._open_folder_button.clicked.connect(self._open_folder)
         frame_layout.addWidget(self._open_folder_button)
-
-        self._create_style_button = QToolButton(self)
-        self._create_style_button.setIcon(Krita.instance().icon("list-add"))
-        self._create_style_button.clicked.connect(self._create_style)
-        frame_layout.addWidget(self._create_style_button)
-
-        self._delete_style_button = QToolButton(self)
-        self._delete_style_button.setIcon(Krita.instance().icon("deletelayer"))
-        self._delete_style_button.clicked.connect(self._delete_style)
-        frame_layout.addWidget(self._delete_style_button)
 
         self._layout.addWidget(frame)
 
@@ -629,8 +629,9 @@ def update_model_lists(self):
         self._read()
 
     def _create_style(self):
+        checkpoint = self._style_widgets["sd_checkpoint"].value
         # make sure the new style is in the combobox before setting it as the current style
-        new_style = Styles.list().create()
+        new_style = Styles.list().create(checkpoint=checkpoint)
         self._update_style_list()
         self.current_style = new_style