Code for V2 added, completely new UI and lots of features

README.md

@@ -1,6 +1,8 @@
 # Diffusion Bee - Stable Diffusion GUI App for MacOS
 [![Twitter](https://img.shields.io/twitter/url.svg?label=Follow%20%40divamgupta&style=social&url=https%3A%2F%2Ftwitter.com%2Fdivamgupta)](https://twitter.com/divamgupta)
 
+a
+
 
 ### Diffusion Bee is the easiest way to run Stable Diffusion locally on your Intel / M1 Mac. Comes with a one-click installer. No dependencies or technical knowledge needed.
 
@@ -14,6 +16,7 @@
 Download at https://diffusionbee.com/
 
 
+
 <br>
 
 For prompt ideas visit https://arthub.ai

backends/model_converter/convert_model.py

@@ -42,6 +42,12 @@ def convert_model(checkpoint_filename=None, out_filename=None,  torch_weights=No
 
 
     for k in (list(state_dict.keys())):
+
+        if ".encoder." in k or ".decoder." in k:
+            for v , s  in [('q' , 'to_q') , ('v' , 'to_v' ) , ('k' , 'to_k') , ('proj_out' , 'to_out.0')]:
+                if f'.{s}.' in k:
+                    state_dict[ k.replace(f'.{s}.' , f'.{v}.') ] = state_dict[k]
+
         if '.norm' in k and '.bias' in k:
             k2 = k.replace(".bias" , ".weight")
             k3 = k.replace(".bias" , ".bias_by_weight")
@@ -106,6 +112,15 @@ def convert_model(checkpoint_filename=None, out_filename=None,  torch_weights=No
         outfile.write_key(key=k , tensor=np_arr)
 
     outfile.finish_write()
+    sd_version = model_type.replace("_float16" , "").replace("_float32" , "") 
+    if sd_version in ["SD_1x" , "SD_2x"]:
+        sd_type = "sd_model"
+    elif sd_version in ["SD_1x_inpaint" , "SD_2x_inpaint"]:
+        sd_type = "sd_model_inpaint"
+    else:
+        raise ValueError("Invalid sd_version "+ sd_version)
+    model_metadata = {"float_type" : cur_dtype , "sd_type" :sd_version, "type" : sd_type }
+    print("__converted_model_data__" , json.dumps(model_metadata))
 
 
 def usage():

backends/model_converter/sd_shapes.py

@@ -77,7 +77,10 @@ def add_aux_shapes(d):
     "SD_1x_float16": 1012 , 
     "SD_1x_inpaint_float16": 1013 , 
     "SD_1x_just_controlnet_16" : 1014,
-    "SD_2x_float16": 1015 }
+    "SD_2x_float16": 1015 ,
+    "sdxl_base_unet_f8_rest_f16" : 3031,
+    }
+
 
 
 extra_keys = ['temb_coefficients_fp32' , 'temb_coefficients_fp16' , 'causal_mask' , 'aux_output_conv.weight' , 'aux_output_conv.bias', 'alphas_cumprod']

backends/model_converter/tdict.py

@@ -22,7 +22,7 @@ def pad_bytes(self, nb):
         self.out_file.write(b'\x00' * nb )
         self.filep += nb
 
-    def write_block(self, block_array):
+    def write_block(self, block_array ):
 
         if self.filep  %64 != 0:
             npad = 64 - self.filep  % 64
@@ -73,7 +73,11 @@ def read_block(self, header_pos, np_shape=None, np_dtype=None):
         block  = self.in_file.read(l)
 
         if np_dtype is not None:
-            block = np.frombuffer(block, dtype=np_dtype)
+            if "cus_" in np_dtype: #custom dtype, which is not supported by numpy 
+                block = np.frombuffer(block, dtype='uint8')
+                block = decode_custom_dtype(block , np_dtype )
+            else:
+                block = np.frombuffer(block, dtype=np_dtype)
 
         if np_shape is not None:
             block = block.reshape(np_shape)
@@ -145,13 +149,24 @@ def read_key(self , key):
         assert len(ret_arr.tobytes()) == w_idx_len
         return ret_arr.copy()
 
-    def write_key(self , key , tensor):
+    def write_key(self , key , tensor, custom_dtype=None ):
         assert key not in self.keys_info
         dtype = str(tensor.dtype)
+        if custom_dtype is not None:
+            dtype = custom_dtype
+            raise not NotImplementedError("Casting not implemented.")
+
         write_info = self.write_block(tensor)
         shape = list(tensor.shape)
         self.keys_info[key] = {"start": write_info['n_start_data'] , "end" : write_info['n_end_data'] , "shape": shape , "dtype" : dtype }
 
+    def write_key_custom_dype(self, uint8_arr , key, custom_dtype , shape ):
+        assert key not in self.keys_info
+        dtype = custom_dtype
+        write_info = self.write_block(uint8_arr)
+        self.keys_info[key] = {"start": write_info['n_start_data'] , "end" : write_info['n_end_data'] , "shape": shape , "dtype" : dtype }
+
+
     def finish_write(self):
 
         if self.filep  %64 != 0:

backends/stable_diffusion/applets/applets.py

@@ -0,0 +1,58 @@
+import json
+
+def update_state_raw(key , val ):
+    print( "utds " +  key + "___U_P_D_A_T_E___" + json.dumps(val))
+
+
+registered_applets = {}
+
+def register_applet(model_container, applet_cls):
+    applet_name = applet_cls.applet_name
+    applet = applet_cls()
+    applet.init_applet(model_container)
+    registered_applets[applet_name] = applet
+
+def run_applet(applet_name , params_dict):
+    registered_applets[applet_name].run(params_dict)
+
+
+class AppletBase:
+
+    applet_name = None
+    applet_title = None
+    is_stop_avail = False
+    applet_description = ""
+    applet_icon = "file"
+    applet_icon_fname = None
+
+    def run(self, params):
+        raise NotImplementedError("base cls")
+
+    def get_input_form(self):
+        return []
+
+    def update_output( self, key , val ):
+        self.update_state( "outputs." + key , val)
+
+    def update_state(self , key  , val ):
+        key = "registered_ext_applets." + self.applet_name + "." +  key
+        update_state_raw(key , val)
+
+    def init_applet(self, model_container):
+
+        self.model_container = model_container
+
+        update_state_raw( "registered_ext_applets." + self.applet_name , {})
+        update_state_raw( "registered_ext_applets." + self.applet_name + ".id" ,  self.applet_name)
+        update_state_raw( "registered_ext_applets." + self.applet_name + ".title" ,  self.applet_title)
+        update_state_raw( "registered_ext_applets." + self.applet_name + ".description" ,  self.applet_description)
+        update_state_raw( "registered_ext_applets." + self.applet_name + ".icon" ,  self.applet_icon )
+        if self.applet_icon_fname is not None:
+            update_state_raw( "registered_ext_applets." + self.applet_name + ".img_icon" ,  self.applet_icon_fname )        
+        update_state_raw( "registered_ext_applets." + self.applet_name + ".home_category" ,  "misc")
+        update_state_raw( "registered_ext_applets." + self.applet_name + ".inputs" ,   self.get_input_form() )
+        update_state_raw( "registered_ext_applets." + self.applet_name + ".outputs" ,   [] )
+        update_state_raw( "registered_ext_applets." + self.applet_name + ".is_stop_avail" ,   self.is_stop_avail )
+
+
+

backends/stable_diffusion/applets/form_utils.py

@@ -0,0 +1,75 @@
+import random
+
+def get_textbox(id, type="str" , default="" ,  title="" , description=""):
+
+    types = {"str" : "text" , "int" : "number" , "float" : "number"}
+
+    return {
+                "id": str(random.randint(0,19989999)),
+                "component": "InputWithDesc",
+                "title": title,
+                "description":description ,
+                "children": [
+                    {
+                        "id": id,
+                        "component": "Textbox",
+                        "placeholder" : "",
+                        "default_value" : default ,
+                        "type" : types[type],
+                        "output_type" :type ,
+                        "is_persistant" : False
+                    }
+                ]
+            }
+
+def get_output_text(text):
+    return {
+        "id":  str(random.randint(0,19989999)),
+        "component": "OutputText",
+        "text" : text
+    }
+def get_output_img(img_path, save_ext='.png' , is_save=False):
+    return {
+        "id":  str(random.randint(0,19989999)),
+        "component": "OutputImage",
+        "img_path" : img_path,
+        "is_save" : is_save , 
+        "save_ext" : save_ext 
+    }
+
+def get_file_textbox(id , path_type="",  title="" , description="" ):
+    return {
+                "id":  str(random.randint(0,19989999)),
+                "component": "InputWithDesc",
+                "full_width": True,
+                "title": title,
+                "description": description,
+                "children": [
+                    {
+                        "id": id ,
+                        "component": "FilePathTextBox",
+                        "placeholder" : "",
+                        "is_persistant" : False,
+                        "path_type" : path_type
+                    },
+                ]
+            }
+
+def get_textarea(id ,   title="" , description="" ):
+    return {
+                "id":  str(random.randint(0,19989999)),
+                "component": "InputWithDesc",
+                "full_width": True,
+                "title": title,
+                "description": description,
+                "children": [
+                    {
+                        "id": id ,
+                        "component": "Textarea",
+                        "placeholder" : title ,
+                        "is_small" : True,
+                        "is_persistant" : False,
+                    },
+                ]
+            }
+