adapter weights for sketch2image

S2I/modules/models.py

@@ -4,7 +4,7 @@ from diffusers import DDPMScheduler
 from transformers import AutoTokenizer, CLIPTextModel
 from diffusers import AutoencoderKL, UNet2DConditionModel
 from peft import LoraConfig
-from S2I.modules.utils import sc_vae_encoder_fwd, sc_vae_decoder_fwd, download_models, get_model_path
+from S2I.modules.utils import sc_vae_encoder_fwd, sc_vae_decoder_fwd, download_models, get_model_path, get_s2i_home
 
 
 class RelationShipConvolution(torch.nn.Module):
@@ -50,7 +50,16 @@ class PrimaryModel:
         vae.decoder.skip_conv_4 = torch.nn.Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False).cuda()
         vae.decoder.ignore_skip = False
         return vae
-
+    def weights_adapter(self, p_ckpt, model_name):
+        if model_name == '350k-adapter':
+            home = get_s2i_home()
+            sd_sketch = torch.load(os.path.join(home, f"sketch2image_lora_350k.pkl"), map_location="cpu")
+            sd = torch.load(p_ckpt, map_location="cpu")
+            sd.update(sd_sketch)
+            return sd 
+        else:
+            sd = torch.load(p_ckpt, map_location="cpu")
+            return sd
     def from_pretrained(self, model_name, r):
         if self.global_tokenizer is None:
             # self.global_tokenizer = AutoTokenizer.from_pretrained(self.backbone_diffusion_path,
@@ -72,7 +81,7 @@ class PrimaryModel:
             self.global_unet = self._load_model(self.backbone_diffusion_path, UNet2DConditionModel, unet_mode=True)
             p_ckpt_path = download_models()
             p_ckpt = get_model_path(model_name=model_name, model_paths=p_ckpt_path)
-            sd = torch.load(p_ckpt, map_location="cpu")
+            sd = self.weights_adapter(p_ckpt, model_name)
             conv_in_pretrained = copy.deepcopy(self.global_unet.conv_in)
             self.global_unet.conv_in = RelationShipConvolution(conv_in_pretrained, self.global_unet.conv_in, r)
             unet_lora_config = LoraConfig(r=sd["rank_unet"], init_lora_weights="gaussian",

S2I/modules/utils.py

@@ -84,6 +84,7 @@ def get_s2i_home() -> str:
 
 def download_models():
     urls = {
+        '350k-adapter': 'https://huggingface.co/myn0908/sk2ks/resolve/main/adapter_weights_large_sketch2image_lora.pkl?download=true',
         '350k': 'https://huggingface.co/myn0908/sk2ks/resolve/main/sketch_to_image_mixed_weights_350k_lora.pkl?download=true',
         '100k': 'https://huggingface.co/myn0908/sk2ks/resolve/main/model_16001.pkl?download=true',
     }

app.py

@@ -263,7 +263,7 @@ with gr.Blocks(css=css, theme="NoCrypt/miku@1.2.1") as demo:
                                         label="Demo Speed",
                                         interactive=True)
                                 model_options = gr.Radio(
-                                        choices=["100k", "350k"],
+                                        choices=["100k", "350k", "350k-adapter"],
                                         value="350k",
                                         label="Type Sketch2Image models",
                                         interactive=True)