Added some more lines to support SD1.x with TINY U-Nets too.

This is a repleacement of my earlier PR leejet#745  (just closed by myself) with much better code.
Also updated doc file about TINY U-Nets.

Author: akleine

README.md

@@ -17,6 +17,7 @@ API and command-line option may change frequently.***
   - Image Models
     - SD1.x, SD2.x, [SD-Turbo](https://huggingface.co/stabilityai/sd-turbo)
     - SDXL, [SDXL-Turbo](https://huggingface.co/stabilityai/sdxl-turbo)
+    - [some SD1.x and SDXL distilled models](./docs/distilled_sd.md)
     - [SD3/SD3.5](./docs/sd3.md)
     - [Flux-dev/Flux-schnell](./docs/flux.md)
     - [Chroma](./docs/chroma.md)

docs/distilled_sd.md

@@ -1,13 +1,15 @@
-# Running distilled SDXL models: SSD1B
+# Running distilled models: SSD1B and SD1.x with tiny U-Nets
 
-### Preface
+## Preface
 
-This kind of models has a reduced U-Net part. Unlike other SDXL models the U-Net has only one middle block and lesser attention layers in up and down blocks, resulting in relatively smaller files. Running these models saves more than 33% of the time. For more details, refer to Segmind's paper on https://arxiv.org/abs/2401.02677v1 .
+This kind of models have a reduced U-Net part. 
+Unlike other SDXL models the U-Net of SSD1B has only one middle block and lesser attention layers in up and down blocks, resulting in relatively smaller files. Running these models saves more than 33% of the time. For more details, refer to Segmind's paper on https://arxiv.org/abs/2401.02677v1 .
+Unlike other SD 1.x models Tiny-UNet models consist of only 6 U-Net blocks, resulting in relatively smaller files (approximately 1 GB). Running these models saves almost 50% of the time. For more details, refer to the paper: https://arxiv.org/pdf/2305.15798.pdf .
 
-### How to Use 
+## SSD1B
 
 Unfortunately not all of this models follow the standard model parameter naming mapping. 
-Anyway there are some useful SSD1B models available online, such as:
+Anyway there are some very useful SSD1B models available online, such as:
 
  * https://huggingface.co/segmind/SSD-1B/resolve/main/SSD-1B-A1111.safetensors
  * https://huggingface.co/hassenhamdi/SSD-1B-fp8_e4m3fn/resolve/main/SSD-1B_fp8_e4m3fn.safetensors 
@@ -16,3 +18,69 @@ Also there are useful LORAs available:
 
  * https://huggingface.co/seungminh/lora-swarovski-SSD-1B/resolve/main/pytorch_lora_weights.safetensors
  * https://huggingface.co/kylielee505/mylcmlorassd/resolve/main/pytorch_lora_weights.safetensors   
+
+You can use this files **out-of-the-box** - unlike models in next section.
+
+
+## SD1.x with tiny U-Nets
+
+There are some Tiny SD 1.x models available online, such as:
+
+ * https://huggingface.co/segmind/tiny-sd
+ * https://huggingface.co/segmind/portrait-finetuned
+ * https://huggingface.co/nota-ai/bk-sdm-tiny
+
+These models need some conversion, for example because partially tensors are **non contiguous** stored. To create a usable checkpoint file, follow these **easy** steps:
+
+### Download model from Hugging Face
+
+Download the model using Python on your computer, for example this way:
+
+```python
+import torch
+from diffusers import StableDiffusionPipeline
+pipe = StableDiffusionPipeline.from_pretrained("segmind/tiny-sd")
+unet=pipe.unet
+for param in unet.parameters():
+    param.data = param.data.contiguous()     # <- important here
+pipe.save_pretrained("segmindtiny-sd", safe_serialization=True)
+```
+
+### Convert that to a ckpt file 
+
+To convert the downloaded model to a checkpoint file, you need another Python script. Download the conversion script from here:
+
+ * https://raw.githubusercontent.com/huggingface/diffusers/refs/heads/main/scripts/convert_diffusers_to_original_stable_diffusion.py
+
+
+### Run convert script
+
+Now, run that conversion script:
+
+```bash
+python convert_diffusers_to_original_stable_diffusion.py \
+	--model_path  ./segmindtiny-sd \
+	--checkpoint_path ./segmind_tiny-sd.ckpt --half
+```
+
+The file **segmind_tiny-sd.ckpt**  will be generated and is now ready to use with sd.cpp
+
+You can follow a similar process for other models mentioned above from Hugging Face. 
+
+
+### Another ckpt file on the net
+
+There is another model file available online: 
+
+ * https://huggingface.co/ClashSAN/small-sd/resolve/main/tinySDdistilled.ckpt
+ 
+If you want to use that, you have to adjust some **non-contiguous tensors** first:
+
+```python
+import torch
+ckpt = torch.load("tinySDdistilled.ckpt", map_location=torch.device('cpu'))
+for key, value in ckpt['state_dict'].items():
+    if isinstance(value, torch.Tensor):
+        ckpt['state_dict'][key] = value.contiguous()
+torch.save(ckpt, "tinySDdistilled_fixed.ckpt")
+```

model.cpp

@@ -623,6 +623,14 @@ std::string convert_tensor_name(std::string name) {
     if (starts_with(name, "diffusion_model")) {
         name = "model." + name;
     }
+    if (starts_with(name, "model.diffusion_model.up_blocks.0.attentions.0.")) {
+        name.replace(0,sizeof("model.diffusion_model.up_blocks.0.attentions.0.") - 1,
+                       "model.diffusion_model.output_blocks.0.1.");
+    }
+    if (starts_with(name, "model.diffusion_model.up_blocks.0.attentions.1.")) {
+        name.replace(0,sizeof("model.diffusion_model.up_blocks.0.attentions.1.") - 1,
+                       "model.diffusion_model.output_blocks.1.1.");
+    }
     // size_t pos = name.find("lora_A");
     // if (pos != std::string::npos) {
     //     name.replace(pos, strlen("lora_A"), "lora_up");
@@ -1887,7 +1895,12 @@ SDVersion ModelLoader::get_sd_version() {
         if (is_ip2p) {
             return VERSION_SD1_PIX2PIX;
         }
-        return VERSION_SD1;
+        for (auto& tensor_storage : tensor_storages) {
+            if (tensor_storage.name.find("model.diffusion_model.middle_block") != std::string::npos) {
+                return VERSION_SD1;      // found a middle block, so it is SD1
+            }
+        }
+        return VERSION_SD1_TINY_UNET;
     } else if (token_embedding_weight.ne[0] == 1024) {
         if (is_inpaint) {
             return VERSION_SD2_INPAINT;

model.h

@@ -22,6 +22,7 @@ enum SDVersion {
     VERSION_SD1,
     VERSION_SD1_INPAINT,
     VERSION_SD1_PIX2PIX,
+    VERSION_SD1_TINY_UNET,
     VERSION_SD2,
     VERSION_SD2_INPAINT,
     VERSION_SDXL,
@@ -42,7 +43,7 @@ enum SDVersion {
 };
 
 static inline bool sd_version_is_sd1(SDVersion version) {
-    if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX) {
+    if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX || version == VERSION_SD1_TINY_UNET) {
         return true;
     }
     return false;

stable-diffusion.cpp

@@ -28,6 +28,7 @@ const char* model_version_to_str[] = {
     "SD 1.x",
     "SD 1.x Inpaint",
     "Instruct-Pix2Pix",
+    "SD 1.x Tiny UNet",
     "SD 2.x",
     "SD 2.x Inpaint",
     "SDXL",

unet.hpp

@@ -204,6 +204,9 @@ class UnetModelBlock : public GGMLBlock {
             adm_in_channels   = 768;
             num_head_channels = 64;
             num_heads         = -1;
+        } else if (version == VERSION_SD1_TINY_UNET) {
+            num_res_blocks    = 1;
+            channel_mult      = {1, 2, 4};
         }
         if (sd_version_is_inpaint(version)) {
             in_channels = 9;
@@ -281,6 +284,9 @@ class UnetModelBlock : public GGMLBlock {
                                                                                       context_dim));
                 }
                 input_block_chans.push_back(ch);
+                if (version == VERSION_SD1_TINY_UNET) {
+                    input_block_idx++;
+                }
             }
             if (i != len_mults - 1) {
                 input_block_idx += 1;
@@ -299,14 +305,16 @@ class UnetModelBlock : public GGMLBlock {
             d_head = num_head_channels;
             n_head = ch / d_head;
         }
-        blocks["middle_block.0"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
-        if (version != VERSION_SDXL_SSD1B) {
-            blocks["middle_block.1"] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
+        if (version != VERSION_SD1_TINY_UNET) {
+            blocks["middle_block.0"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
+            if (version != VERSION_SDXL_SSD1B) {
+                blocks["middle_block.1"] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
                                                                                       n_head,
                                                                                       d_head,
                                                                                       transformer_depth[transformer_depth.size() - 1],
                                                                                       context_dim));
-            blocks["middle_block.2"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
+                blocks["middle_block.2"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
+            }
         }
         // output_blocks
         int output_block_idx = 0;
@@ -340,6 +348,12 @@ class UnetModelBlock : public GGMLBlock {
                 }
 
                 if (i > 0 && j == num_res_blocks) {
+                    if (version == VERSION_SD1_TINY_UNET) {
+                        output_block_idx++;
+                        if (output_block_idx == 2) {
+                            up_sample_idx=1;
+                        }
+                    }
                     std::string name = "output_blocks." + std::to_string(output_block_idx) + "." + std::to_string(up_sample_idx);
                     blocks[name]     = std::shared_ptr<GGMLBlock>(new UpSampleBlock(ch, ch));
 
@@ -473,6 +487,9 @@ class UnetModelBlock : public GGMLBlock {
                 }
                 hs.push_back(h);
             }
+            if (version == VERSION_SD1_TINY_UNET) {
+                input_block_idx++;
+            }
             if (i != len_mults - 1) {
                 ds *= 2;
                 input_block_idx += 1;
@@ -487,10 +504,12 @@ class UnetModelBlock : public GGMLBlock {
         // [N, 4*model_channels, h/8, w/8]
 
         // middle_block
-        h = resblock_forward("middle_block.0", ctx, h, emb, num_video_frames);                      // [N, 4*model_channels, h/8, w/8]
-        if (version != VERSION_SDXL_SSD1B) {
-            h = attention_layer_forward("middle_block.1", ctx, backend, h, context, num_video_frames);  // [N, 4*model_channels, h/8, w/8]
-            h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames);                      // [N, 4*model_channels, h/8, w/8]
+        if (version != VERSION_SD1_TINY_UNET) {
+            h = resblock_forward("middle_block.0", ctx, h, emb, num_video_frames);                      // [N, 4*model_channels, h/8, w/8]
+            if (version != VERSION_SDXL_SSD1B) {
+                h = attention_layer_forward("middle_block.1", ctx, backend, h, context, num_video_frames);  // [N, 4*model_channels, h/8, w/8]
+                h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames);                      // [N, 4*model_channels, h/8, w/8]
+            }
         }
         if (controls.size() > 0) {
             auto cs = ggml_scale_inplace(ctx, controls[controls.size() - 1], control_strength);
@@ -527,6 +546,12 @@ class UnetModelBlock : public GGMLBlock {
                 }
 
                 if (i > 0 && j == num_res_blocks) {
+                    if (version == VERSION_SD1_TINY_UNET) {
+                        output_block_idx++;
+                        if (output_block_idx == 2) {
+                            up_sample_idx=1;
+                        }
+                    }
                     std::string name = "output_blocks." + std::to_string(output_block_idx) + "." + std::to_string(up_sample_idx);
                     auto block       = std::dynamic_pointer_cast<UpSampleBlock>(blocks[name]);