[DEBUG] Dedug custom nn

app.py

@@ -41,6 +41,7 @@ def load_pipeline(model_id,
     if model_id in model_cache:
         return model_cache[model_id]
 
+
     if controlnet_checkbox:
         if controlnet_mode == "depth_map":
             controlnet = ControlNetModel.from_pretrained(
@@ -72,39 +73,58 @@ def load_pipeline(model_id,
                 cache_dir="./models_cache",
                 torch_dtype=torch_dtype
             )
-        pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id,
-                                                                 controlnet=controlnet,
-                                                                 torch_dtype=torch_dtype,
-                                                                 safety_checker=None).to(device)
+        if model_id == "YaArtemNosenko/dino_stickers":
+            # Use the specified base model for your LoRA adapter.
+            base_model = "CompVis/stable-diffusion-v1-4"
+            # Load the LoRA weights
+            pipe = StableDiffusionControlNetPipeline.from_pretrained(base_model, 
+                                                                     controlnet=controlnet,
+                                                                     torch_dtype=torch_dtype, 
+                                                                     safety_checker=None).to(device)
+            pipe.unet = PeftModel.from_pretrained(
+                pipe.unet, 
+                model_id, 
+                subfolder="unet", 
+                torch_dtype=torch_dtype
+            )
+            pipe.text_encoder = PeftModel.from_pretrained(
+                pipe.text_encoder, 
+                model_id, 
+                subfolder="text_encoder", 
+                torch_dtype=torch_dtype
+            )
+        else:
+            pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id,
+                                                                     controlnet=controlnet,
+                                                                     torch_dtype=torch_dtype,
+                                                                     safety_checker=None).to(device)
         # params['image'] = controlnet_image
         # params['controlnet_conditioning_scale'] = float(controlnet_strength)
     else:
-        pipe = StableDiffusionPipeline.from_pretrained(model_id,
-                                                       torch_dtype=torch_dtype,
-                                                       safety_checker=None).to(device)
-
-    if model_id == "YaArtemNosenko/dino_stickers":
-        # Use the specified base model for your LoRA adapter.
-        base_model = "CompVis/stable-diffusion-v1-4"
-        # Load the LoRA weights
-        pipe.unet = PeftModel.from_pretrained(
-            pipe.unet, 
-            model_id, 
-            subfolder="unet", 
-            torch_dtype=torch_dtype
-        )
-        pipe.text_encoder = PeftModel.from_pretrained(
-            pipe.text_encoder, 
-            model_id, 
-            subfolder="text_encoder", 
-            torch_dtype=torch_dtype
-        )
-        pipe.unet.load_state_dict({k: lora_scale * v for k, v in pipe.unet.state_dict().items()})
-        pipe.text_encoder.load_state_dict({k: lora_scale * v for k, v in pipe.text_encoder.state_dict().items()})
-    else:
-        pipe = StableDiffusionPipeline.from_pretrained(model_id,
-                                                 torch_dtype=torch_dtype
-                                                 )
+        if model_id == "YaArtemNosenko/dino_stickers":
+            base_model = "CompVis/stable-diffusion-v1-4"
+            pipe = StableDiffusionPipeline.from_pretrained(base_model, torch_dtype=torch_dtype)
+            # Load the LoRA weights
+            pipe.unet = PeftModel.from_pretrained(
+                                                pipe.unet, 
+                                                model_id, 
+                                                subfolder="unet", 
+                                                torch_dtype=torch_dtype
+            )
+
+            pipe.text_encoder = PeftModel.from_pretrained(
+                pipe.text_encoder, 
+                model_id, 
+                subfolder="text_encoder", 
+                torch_dtype=torch_dtype
+            )
+        else:
+            pipe = StableDiffusionPipeline.from_pretrained(model_id,
+                                                           torch_dtype=torch_dtype,
+                                                           safety_checker=None).to(device)
+
+            pipe.unet.load_state_dict({k: lora_scale * v for k, v in pipe.unet.state_dict().items()})
+            pipe.text_encoder.load_state_dict({k: lora_scale * v for k, v in pipe.text_encoder.state_dict().items()})
 
     if ip_adapter_checkbox:
         pipe.load_ip_adapter("h94/IP-Adapter",