├── README.md ├── checkpoint_merging └── merge_checkpoints.py ├── helpers ├── file_utils.py └── tensor_utils.py ├── hindi guide └── chapter1.md ├── lora_merging ├── merge_lora_checkpoint.py └── merge_two_lora.py ├── main.py └── requirements.txt /README.md: -------------------------------------------------------------------------------- 1 | # Merge Diffusion Tool 2 | 3 | **Merge Diffusion Tool** is an open-source solution developed by **EnhanceAI.art**, providing seamless ways to blend LoRA models, integrate LoRA into checkpoints, and merge Stable Diffusion checkpoints. Enhance your AI workflows with this powerful merging tool, designed to support a wide range of diffusion models like **Flux Dev**, **Flux Schnell**, **Stable Diffusion 1.5**, **SD2**, **SD3**, and **SDXL**. 4 | 5 | Discover more advanced AI tools at [EnhanceAI.art](https://enhanceai.art). 6 | 7 | ## Features 8 | 9 | - **Merge Two LoRA Models** with adjustable blending ratios. 10 | - **Integrate LoRA into a Checkpoint** for enhanced model performance. 11 | - **Merge Two Checkpoints** using custom blend ratios. 12 | - Full support for `.safetensors` format, ensuring efficient and safe handling. 13 | 14 | ### Key Advantages: 15 | - Streamline your AI model merging process. 16 | - Built with flexibility to cater to diverse AI model formats. 17 | - Completely open-source with community support via Discord. 18 | 19 | --- 20 | 21 | ## Installation 22 | 23 | To set up the tool, follow these steps: 24 | 25 | 1. Clone the repository: 26 | ```bash 27 | git clone https://github.com/itspranavajay/Merge-Diffusion-Tool.git 28 | cd merge-diffusion-tool 29 | ``` 30 | 31 | 2. Install the required dependencies: 32 | ```bash 33 | pip install -r requirements.txt 34 | ``` 35 | 36 | ### Dependencies: 37 | 38 | - **Python 3.8+** 39 | - **PyTorch** for deep learning operations. 40 | - **Safetensors** library for model file handling. 41 | 42 | --- 43 | 44 | ## Usage 45 | 46 | The **Merge Diffusion Tool** provides three main operations: 47 | 48 | 1. **Merge Two LoRA Models** 49 | 2. **Merge LoRA into a Checkpoint** 50 | 3. **Merge Two Checkpoints** 51 | 52 | ### Command-Line Arguments: 53 | 54 | | Argument | Description | 55 | |------------------|-----------------------------------------------------------------------------| 56 | | `--operation` | Select operation: `merge_loras`, `merge_lora_checkpoint`, `merge_checkpoints`| 57 | | `--model1` | Path to the first model file. | 58 | | `--model2` | Path to the second model file (optional for LoRA into checkpoint). | 59 | | `--output` | Output path for the merged model. | 60 | | `--alpha` | Blend ratio for merging models (default: 0.5). | 61 | | `--merge_weight` | Weight for merging LoRA into checkpoint (default: 0.5). | 62 | 63 | ### Example 1: Merging Two LoRA Models 64 | 65 | ```bash 66 | python main.py --operation merge_loras --model1 lora1.safetensors --model2 lora2.safetensors --output output_lora.safetensors --alpha 0.7 67 | ``` 68 | 69 | This command merges `lora1.safetensors` and `lora2.safetensors`, with 70% contribution from `lora1`. 70 | 71 | ### Example 2: Merging LoRA into a Checkpoint 72 | 73 | ```bash 74 | python main.py --operation merge_lora_checkpoint --model1 lora_model.safetensors --model2 checkpoint_model.safetensors --output output_checkpoint.safetensors --merge_weight 0.6 75 | ``` 76 | 77 | In this example, the LoRA model merges into the checkpoint with 60% influence. 78 | 79 | ### Example 3: Merging Two Checkpoints 80 | 81 | ```bash 82 | python main.py --operation merge_checkpoints --model1 checkpoint1.safetensors --model2 checkpoint2.safetensors --output output_checkpoint.safetensors --alpha 0.5 83 | ``` 84 | 85 | Merges both checkpoint models in a 50-50 blend. 86 | 87 | --- 88 | 89 | ## Supported Formats 90 | 91 | - **`.safetensors`**: Ensures safe, optimized storage and handling of AI models. 92 | 93 | --- 94 | 95 | ## Donation EnhanceAI.art 96 | 97 | 👏 If you find this tool helpful, please consider supporting the development by checking out our pricing plans at [enhanceai.art/pricing](https://enhanceai.art/pricing). By purchasing any plan, not only will you be supporting future open-source projects like this, but you'll also unlock **many exclusive AI features** that will greatly enhance yourself AI workflows! 98 | 99 | --- 100 | 101 | ## Support 102 | 103 | For assistance, join our [Discord server](https://discord.gg/wtyHstGZ) to connect with the community and developers. 104 | 105 | --- 106 | 107 | ## Contributing 108 | 109 | We welcome contributions! Feel free to open issues, submit pull requests, or suggest features that would improve this tool. All contributions are highly appreciated. 110 | -------------------------------------------------------------------------------- /checkpoint_merging/merge_checkpoints.py: -------------------------------------------------------------------------------- 1 | from tqdm import tqdm 2 | from helpers.tensor_utils import resize_tensor_shapes 3 | 4 | def merge_checkpoints(ckpt1, ckpt2, blend_ratio=0.5): 5 | print(f"Merging checkpoints with blend ratio: {blend_ratio}") 6 | merged = {} 7 | all_keys = set(ckpt1.keys()).union(set(ckpt2.keys())) 8 | 9 | for key in tqdm(all_keys, desc="Merging Checkpoints", unit="layer"): 10 | t1, t2 = ckpt1.get(key), ckpt2.get(key) 11 | if t1 is not None and t2 is not None: 12 | t1, t2 = resize_tensor_shapes(t1, t2) 13 | merged[key] = blend_ratio * t1 + (1 - blend_ratio) * t2 14 | elif t1 is not None: 15 | merged[key] = t1 16 | else: 17 | merged[key] = t2 18 | 19 | return merged 20 | -------------------------------------------------------------------------------- /helpers/file_utils.py: -------------------------------------------------------------------------------- 1 | from safetensors.torch import load_file, save_file 2 | 3 | def load_model(file_path): 4 | return load_file(file_path) 5 | 6 | def save_model(merged_model, output_file): 7 | print(f"Saving merged model to {output_file}") 8 | save_file(merged_model, output_file) 9 | -------------------------------------------------------------------------------- /helpers/tensor_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | def resize_tensor_shapes(tensor1, tensor2): 5 | if tensor1.size() == tensor2.size(): 6 | return tensor1, tensor2 7 | 8 | max_shape = [max(s1, s2) for s1, s2 in zip(tensor1.shape, tensor2.shape)] 9 | tensor1_resized = F.pad(tensor1, (0, max_shape[-1] - tensor1.size(-1))) 10 | tensor2_resized = F.pad(tensor2, (0, max_shape[-1] - tensor2.size(-1))) 11 | 12 | return tensor1_resized, tensor2_resized 13 | -------------------------------------------------------------------------------- /hindi guide/chapter1.md: -------------------------------------------------------------------------------- 1 | ### **Dreambooth Ka Kamaal: Ek Mazedaar Guide to Master Stable Diffusion** 2 | 3 | --- 4 | 5 | #### **Chapter 1: Basics – Dreambooth Mein Kadam Rakhna!** 6 | 7 | Kahani shuru hoti hai ek aise adventure se jahan tum AI ka ek naya superpower, Dreambooth, seekhne wale ho. Lekin, pehle thoda warm-up toh banta hai. Basics samjhe bina, Stable Diffusion ki duniya mein ghusna bilkul waisa hi hai jaise bina handle ke cycle chalana. So, let’s get started! 8 | 9 | --- 10 | 11 | ##### **Weight, Layer aur Bias Ka Logic:** 12 | 13 | Weight, layers, aur bias ko samajhna koi rocket science nahin, bas ek thoda sa patience chahiye. Jab AI training kar raha hota hai, uska ek simple motto hota hai – "Yeh sahi kaam kaise karoon?" Iska jawab hota hai **weight** aur **bias** adjust karna. 14 | 15 | - **Weight:** AI ke andar ek mathematical operation hota hai, jahan input data ko kuch weights diye jaate hain. Yeh weights decide karte hain ki output kaisa hoga. Basically, input ka importance kis weight pe dependent hota hai. 16 | 17 | - **Layer:** Tumhe har neural network ek layer-based structure mein milta hai. Har layer input data ko thoda thoda modify karti hai, aur next layer tak pass karti hai. Aise samjho ki yeh AI ka cooking process hai – ek layer salt dalti hai, doosri layer thoda masala, aur final dish tumhare saamne AI serve karta hai. 18 | 19 | - **Bias:** Bias wo hai jo result ko thoda tilt karta hai. Socho, tumhe ek fixed preference dena hai. Tumhara AI bhi kabhi kabhi koi preferred answer pe focus karne lagta hai, aur yeh bias ka kaam hai. 20 | 21 | Chalo, ab isko code ke zariye samjhte hain: 22 | 23 | ```python 24 | import torch 25 | import torch.nn as nn 26 | 27 | # Ek simple neural network model ka structure 28 | class SimpleModel(nn.Module): 29 | def __init__(self): 30 | super(SimpleModel, self).__init__() 31 | # Define layers: input layer, hidden layer, output layer 32 | self.input_layer = nn.Linear(10, 50) # 10 input features, 50 neurons in hidden layer 33 | self.hidden_layer = nn.Linear(50, 50) # 50 neurons in, 50 neurons out 34 | self.output_layer = nn.Linear(50, 1) # Final output (e.g., regression task) 35 | 36 | # Activation function - sigmoid 37 | self.activation = nn.Sigmoid() 38 | 39 | def forward(self, x): 40 | # Forward pass - weights aur bias adjust ho rahe hain 41 | x = self.activation(self.input_layer(x)) # Input -> Hidden Layer 42 | x = self.activation(self.hidden_layer(x)) # Hidden -> Hidden 43 | x = self.output_layer(x) # Final output 44 | return x 45 | 46 | # Model initialization 47 | model = SimpleModel() 48 | 49 | # Random data ke saath forward pass dekhte hain 50 | input_data = torch.rand(1, 10) # 1 sample, 10 features 51 | output = model(input_data) 52 | print(output) 53 | ``` 54 | 55 | --- 56 | 57 | ##### **Finetuning: Aise Karo AI Ko Master** 58 | 59 | Ab samjho, tumhara AI model basic training ke baad gym join karna chahta hai – **Finetuning** wahi gym hai jahan AI apne pehle se trained knowledge ko aur perfect banata hai. 60 | 61 | Finetuning ka logic kuch aisa hota hai: 62 | 63 | 1. **Pre-trained model** ko use karo – jo already seekh chuka hai kuch important patterns. 64 | 2. Us model ke kuch layers ko “freeze” kar do, taaki uske kuch pre-learned weights change na hon. 65 | 3. **Unfreeze karo** sirf final layers, aur new data ke hisaab se model ko train karo. 66 | 67 | Isko hum ek example se samjhte hain. Socho tum ek Stable Diffusion model ko Dreambooth ke through finetune karna chahte ho. Sabse pehle, pre-trained model ko load karte hain aur sirf last layers ko finetune karte hain: 68 | 69 | ```python 70 | from transformers import StableDiffusionModel 71 | 72 | # Pre-trained model load kar rahe hain 73 | model = StableDiffusionModel.from_pretrained('CompVis/stable-diffusion-v1-4') 74 | 75 | # Sabse pehle model ke layers ko freeze karte hain 76 | for param in model.parameters(): 77 | param.requires_grad = False # Freeze all layers 78 | 79 | # Final layers ko finetuning ke liye unlock karte hain 80 | for param in model.output_layer.parameters(): 81 | param.requires_grad = True # Unfreeze last layers for fine-tuning 82 | 83 | # Now, ab hum nayi data ke saath finetuning karenge 84 | ``` 85 | 86 | --- 87 | 88 | ##### **LoRA Training: Yeh Kya Hai?** 89 | 90 | Ab tumne basic samajh liya, lekin ek new trick seekhne ka time hai: **LoRA (Low-Rank Adaptation)**. Iska fayda yeh hai ki tum model ko finetune kar sakte ho bina saari memory use kiye hue. Yeh bilkul waise hai jaise tum apne memory card pe only **important files** rakhte ho, bakiyon ko ignore karte ho. 91 | 92 | LoRA model training ka core logic yeh hota hai ki instead of training full weights, hum sirf kuch low-rank matrices ko train karte hain, jo computational cost aur memory ko kam kar deta hai. 93 | 94 | --- 95 | 96 | #### **Aur Phir Humare Hero Ka Pehla Kadam Tha!** 97 | 98 | Toh bas, tumne pehla chapter cover kar liya! Tumhe ab Stable Diffusion ke basics, weight aur bias ka kaam, aur finetuning ka logic samajh aa gaya hai. Agle chapter mein hum dekhenge LoRA training ka in-depth process aur model merging ka mazedaar funda. 99 | 100 | Stay tuned for **Chapter 2: Finetuning Ka Rasgulla!** 101 | 102 | -------------------------------------------------------------------------------- /lora_merging/merge_lora_checkpoint.py: -------------------------------------------------------------------------------- 1 | from tqdm import tqdm 2 | from helpers.tensor_utils import resize_tensor_shapes 3 | 4 | def merge_lora_into_checkpoint(lora, checkpoint, merge_weight=0.5): 5 | print(f"Merging LoRA into checkpoint with weight: {merge_weight}") 6 | merged = {} 7 | all_keys = set(checkpoint.keys()).union(set(lora.keys())) 8 | 9 | for key in tqdm(all_keys, desc="Merging LoRA into Checkpoint", unit="layer"): 10 | ckpt_tensor, lora_tensor = checkpoint.get(key), lora.get(key) 11 | if ckpt_tensor is not None and lora_tensor is not None: 12 | ckpt_tensor, lora_tensor = resize_tensor_shapes(ckpt_tensor, lora_tensor) 13 | merged[key] = ckpt_tensor + merge_weight * lora_tensor 14 | elif ckpt_tensor is not None: 15 | merged[key] = ckpt_tensor 16 | else: 17 | merged[key] = merge_weight * lora_tensor 18 | 19 | return merged 20 | -------------------------------------------------------------------------------- /lora_merging/merge_two_lora.py: -------------------------------------------------------------------------------- 1 | from tqdm import tqdm 2 | from helpers.tensor_utils import resize_tensor_shapes 3 | 4 | def merge_lora_models(lora1, lora2, blend_ratio=0.5): 5 | print(f"Merging LoRA models with blend ratio: {blend_ratio}") 6 | merged = {} 7 | all_keys = set(lora1.keys()).union(set(lora2.keys())) 8 | 9 | for key in tqdm(all_keys, desc="Merging LoRA", unit="layer"): 10 | t1, t2 = lora1.get(key), lora2.get(key) 11 | if t1 is not None and t2 is not None: 12 | t1, t2 = resize_tensor_shapes(t1, t2) 13 | merged[key] = blend_ratio * t1 + (1 - blend_ratio) * t2 14 | elif t1 is not None: 15 | merged[key] = t1 16 | else: 17 | merged[key] = t2 18 | 19 | return merged 20 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | from lora_merging.merge_two_lora import merge_lora_models 2 | from lora_merging.merge_lora_checkpoint import merge_lora_into_checkpoint 3 | from checkpoint_merging.merge_checkpoints import merge_checkpoints 4 | from helpers.file_utils import load_model, save_model 5 | 6 | def main(): 7 | import argparse 8 | parser = argparse.ArgumentParser(description="Merging LoRA and Checkpoint Models") 9 | parser.add_argument("--operation", type=str, required=True, choices=["merge_loras", "merge_lora_checkpoint", "merge_checkpoints"], help="Choose operation") 10 | parser.add_argument("--model1", type=str, required=True, help="Path to the first model file") 11 | parser.add_argument("--model2", type=str, required=False, help="Path to the second model file") 12 | parser.add_argument("--output", type=str, required=True, help="Path for the output merged model") 13 | parser.add_argument("--alpha", type=float, default=0.5, help="Blend ratio for merging two models") 14 | parser.add_argument("--merge_weight", type=float, default=0.5, help="Weight for merging LoRA into a checkpoint") 15 | 16 | args = parser.parse_args() 17 | 18 | # Load models 19 | model1 = load_model(args.model1) 20 | model2 = load_model(args.model2) if args.model2 else None 21 | 22 | # Perform merging operations 23 | if args.operation == "merge_loras": 24 | if model2 is None: 25 | print("Error: Please provide two LoRA models for merging.") 26 | return 27 | merged_model = merge_lora_models(model1, model2, args.alpha) 28 | 29 | elif args.operation == "merge_lora_checkpoint": 30 | if model2 is None: 31 | print("Error: Please provide a checkpoint to merge the LoRA into.") 32 | return 33 | merged_model = merge_lora_into_checkpoint(model1, model2, args.merge_weight) 34 | 35 | elif args.operation == "merge_checkpoints": 36 | if model2 is None: 37 | print("Error: Please provide two checkpoints for merging.") 38 | return 39 | merged_model = merge_checkpoints(model1, model2, args.alpha) 40 | 41 | # Save the merged model 42 | save_model(merged_model, args.output) 43 | 44 | if __name__ == "__main__": 45 | main() 46 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch 2 | tqdm 3 | safetensors 4 | --------------------------------------------------------------------------------