From 2825955bcfa30ad2c1d164ca33317ebada544b45 Mon Sep 17 00:00:00 2001
From: Mackerel <mackerel@chickatrice.com>
Date: Wed, 5 Jul 2023 22:18:31 -0400
Subject: [PATCH] add filename (not prefix) option and update readme

---
 README.md   | 17 ++++++++++-------
 __init__.py | 26 ++++++++++++++++++--------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index ff8d628..ffd6f0c 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
 Auto-MBW for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) loosely based on [sdweb-auto-MBW](https://github.com/Xerxemi/sdweb-auto-MBW)
 
 ### Purpose
-This node "advanced > auto merge block weighted" takes two models, merges individual blocks together at various ratios, and automatically rates each merge, keeping the ratio with the highest score. Whether this is a good idea or not is anyone's guess. In practice this makes models that make images the classifier says are good.
+This node "advanced > auto merge block weighted" takes two models, merges individual blocks together at various ratios, and automatically rates each merge, keeping the ratio with the highest score. Whether this is a good idea or not is anyone's guess. In practice this makes models that make images the classifier says are good. You would probably disagree with the classifiers' decisions often.
 
 ### Settings
 - Prompt: to generate sample images to be rated
@@ -10,10 +10,10 @@ This node "advanced > auto merge block weighted" takes two models, merges indivi
 - Classifier: model used to rate images
 
 ### Search Depth
-To calculate ratios to test, the node branches out from powers of 0.5
+To calculate ratios to test, the node branches out from powers of 0.5 
 
-- A depth of 2 will examine 0.0, 0.5, 1.0
-- A depth of 4 will examine 0.0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0
+- A depth of 2 will examine 0.0, 0.5, 1.0 
+- A depth of 4 will examine 0.0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0 
 - A depth of 6 will examine 33 different ratios
 
 There are 25 blocks to examine. If you use a depth of 4 and create 2 samples each, `25 * 9 * 2 = 450` images will be generated.
@@ -26,10 +26,13 @@ The classifier models have been taken from the sdweb-auto-MBW repo.
 - [Cafe Waifu](https://huggingface.co/cafeai/cafe_waifu) and [Cafe Aesthetic](https://huggingface.co/cafeai/cafe_aesthetic)
 
 ### Notes
-- --highvram flag recommended - both models will be kept in VRAM and the process is much faster
-- many hardcoded settings are arbitrary - such as the sampler and block processing order
+- many hardcoded settings are arbitrary such as the seed, sampler and block processing order
 - generated images are not saved
-- the final model is saved in the models/checkpoints directory with a timestamped name
 - the resulting model will contain the text encoder and VAE sent to the node
+
+### Bugs
+- filename box doesn't use the standard comfy "prefix" method
+- merging process doesn't use the comfy ModelPatcher method and takes hundreds of milliseconds
+- - as a result, --highvram flag recommended. both models will be kept in VRAM and the process is much faster
 - the unet will (probably) be fp16 and the rest fp32. that's how they're sent to the node
 - - see: `model_management.should_use_fp16()`
diff --git a/__init__.py b/__init__.py
index ef0e9f9..9879bdb 100644
--- a/__init__.py
+++ b/__init__.py
@@ -2,7 +2,6 @@ import importlib
 import math
 import pathlib
 import sys
-import time
 import warnings
 
 import numpy as np
@@ -43,6 +42,7 @@ class AutoMBW:
                 "search_depth": ("INT", {"default": 4, "min": 2}),
                 "sample_count": ("INT", {"default": 1, "min": 1}),
                 "classifier": (classifiers.__all__,),
+                "filename": ("STRING", { "multiline": False, "default": "ambw" }),
             }}
 
     RETURN_TYPES = ()
@@ -103,7 +103,7 @@ class AutoMBW:
         return maximum
 
     def ambw(self, model1, model2, clip, vae, prompt, negative, search_depth,
-             sample_count, classifier):
+             sample_count, classifier, filename):
         # python setup
         self.model1 = model1
         self.model2 = model2
@@ -150,14 +150,24 @@ class AutoMBW:
         for key in clip:
             sd1[f"cond_stage_model.{key}"] = clip[key]
 
+        create_ckpt = False
+        if filename.endswith(".safetensors"):
+            filename = filename[0:-12]
+        elif filename.endswith(".ckpt"):
+            filename = filename[0:-5]
+            create_ckpt = True
+
         filename = pathlib.Path(folder_paths.folder_names_and_paths[
-            "checkpoints"][0][0]).joinpath(f"ambw{int(time.time())}")
+            "checkpoints"][0][0]).joinpath(f"{filename}")
         print(f"saving as {filename}", end="")
-        try:
-            import safetensors.torch
-            print(".safetensors")
-            safetensors.torch.save_file(sd1, f"{filename}.safetensors")
-        except ModuleNotFoundError:
+        if not create_ckpt:
+            try:
+                import safetensors.torch
+                print(".safetensors")
+                safetensors.torch.save_file(sd1, f"{filename}.safetensors")
+            except ModuleNotFoundError:
+                create_ckpt = True
+        if create_ckpt:
             print(".ckpt")
             torch.save(sd1, f"{filename}.ckpt")