Upload 3 files

blora_utils.py

@@ -0,0 +1,46 @@
+from typing import Optional
+
+BLOCKS = {
+    'content': ['unet.up_blocks.0.attentions.0'],
+    'style': ['unet.up_blocks.0.attentions.1'],
+}
+
+
+def is_belong_to_blocks(key, blocks):
+    try:
+        for g in blocks:
+            if g in key:
+                return True
+        return False
+    except Exception as e:
+        raise type(e)(f'failed to is_belong_to_block, due to: {e}')
+
+
+def filter_lora(state_dict, blocks_):
+    try:
+        return {k: v for k, v in state_dict.items() if is_belong_to_blocks(k, blocks_)}
+    except Exception as e:
+        raise type(e)(f'failed to filter_lora, due to: {e}')
+
+
+def scale_lora(state_dict, alpha):
+    try:
+        return {k: v * alpha for k, v in state_dict.items()}
+    except Exception as e:
+        raise type(e)(f'failed to scale_lora, due to: {e}')
+
+
+def get_target_modules(unet, blocks=None):
+    try:
+        if not blocks:
+            blocks = [('.').join(blk.split('.')[1:]) for blk in BLOCKS['content'] + BLOCKS['style']]
+
+        attns = [attn_processor_name.rsplit('.', 1)[0] for attn_processor_name, _ in unet.attn_processors.items() if
+                 is_belong_to_blocks(attn_processor_name, blocks)]
+
+        target_modules = [f'{attn}.{mat}' for mat in ["to_k", "to_q", "to_v", "to_out.0"] for attn in attns]
+        return target_modules
+    except Exception as e:
+        raise type(e)(f'failed to get_target_modules, due to: {e}')
+
+

inference.py

@@ -0,0 +1,69 @@
+import argparse
+
+import torch
+from diffusers import StableDiffusionXLPipeline, AutoencoderKL
+
+from blora_utils import BLOCKS, filter_lora, scale_lora
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--prompt", type=str, required=True, help="B-LoRA prompt"
+    )
+    parser.add_argument(
+        "--output_path", type=str, required=True, help="path to save the images"
+    )
+    parser.add_argument(
+        "--content_B_LoRA", type=str, default=None, help="path for the content B-LoRA"
+    )
+    parser.add_argument(
+        "--style_B_LoRA", type=str, default=None, help="path for the style B-LoRA"
+    )
+    parser.add_argument(
+        "--content_alpha", type=float, default=1., help="alpha parameter to scale the content B-LoRA weights"
+    )
+    parser.add_argument(
+        "--style_alpha", type=float, default=1., help="alpha parameter to scale the style B-LoRA weights"
+    )
+    parser.add_argument(
+        "--num_images_per_prompt", type=int, default=4, help="number of images per prompt"
+    )
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+    pipeline = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
+                                                         vae=vae,
+                                                         torch_dtype=torch.float16).to("cuda")
+
+    # Get Content B-LoRA SD
+    if args.content_B_LoRA is not None:
+        content_B_LoRA_sd, _ = pipeline.lora_state_dict(args.content_B_LoRA)
+        content_B_LoRA = filter_lora(content_B_LoRA_sd, BLOCKS['content'])
+        content_B_LoRA = scale_lora(content_B_LoRA, args.content_alpha)
+    else:
+        content_B_LoRA = {}
+
+    # Get Style B-LoRA SD
+    if args.style_B_LoRA is not None:
+        style_B_LoRA_sd, _ = pipeline.lora_state_dict(args.style_B_LoRA)
+        style_B_LoRA = filter_lora(style_B_LoRA_sd, BLOCKS['style'])
+        style_B_LoRA = scale_lora(style_B_LoRA, args.style_alpha)
+    else:
+        style_B_LoRA = {}
+
+    # Merge B-LoRAs SD
+    res_lora = {**content_B_LoRA, **style_B_LoRA}
+
+    # Load
+    pipeline.load_lora_into_unet(res_lora, None, pipeline.unet)
+
+    # Generate
+    images = pipeline(args.prompt, num_images_per_prompt=args.num_images_per_prompt).images
+
+    # Save
+    for i, img in enumerate(images):
+        img.save(f'{args.output_path}/{args.prompt}_{i}.jpg')

requirements.txt

@@ -0,0 +1,11 @@
+accelerate
+bitsandbytes==0.36.0.post2
+datasets
+diffusers==0.25.0
+ftfy==6.1.1
+huggingface-hub
+Pillow==9.4.0
+python-slugify==7.0.0
+torch
+torchvision
+transformers