Update processing/setup.py

processing/setup.py

@@ -1,21 +1,20 @@
 import huggingface_hub
 import torch
-from diffusers import ControlNetModel, StableDiffusionXLControlNetInpaintPipeline, AutoencoderKL
+from diffusers import ControlNetModel, StableDiffusionXLControlNetInpaintPipeline
+
 from DPT.dpt.models import DPTDepthModel
 from ip_adapter import IPAdapter, IPAdapterXL
 from ip_adapter.utils import register_cross_attention_hook
 
+
 def setup(base_model_path="stabilityai/stable-diffusion-xl-base-1.0",
           image_encoder_path="sdxl_models/image_encoder",
           ip_ckpt="sdxl_models/ip-adapter_sdxl.bin",
           controlnet_path="diffusers/controlnet-depth-sdxl-1.0",
-          lora_model_path="PixelArt_v1_LoRA_XL",  # Path to your LoRA model
           device="cuda",
           model_depth_path="DPT/weights/dpt_hybrid-midas-501f0c75.pt",
           depth_backbone="vitb_rn50_384"):
-    """Set up the processing module with LoRA, ControlNet, IP Adapter, and Depth Model."""
-    
-    # Ensure that the necessary files are downloaded locally
+    """Set up the processing module."""
     huggingface_hub.snapshot_download(
         repo_id='h94/IP-Adapter',
         allow_patterns=[
@@ -25,17 +24,12 @@ def setup(base_model_path="stabilityai/stable-diffusion-xl-base-1.0",
         local_dir='./',
         local_dir_use_symlinks=False,
     )
-    
-    torch.cuda.empty_cache()
 
-    # Load ControlNet model
-    controlnet = ControlNetModel.from_pretrained(
-        controlnet_path, 
-        use_safetensors=True,
-        torch_dtype=torch.float16
-    ).to(device)
+    torch.cuda.empty_cache()
 
-    # Load Stable Diffusion XL pipeline with ControlNet and inpainting
+    # load SDXL pipeline
+    controlnet = ControlNetModel.from_pretrained(controlnet_path, use_safetensors=True,
+                                                 torch_dtype=torch.float16).to(device)
     pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(
         base_model_path,
         controlnet=controlnet,
@@ -43,31 +37,21 @@ def setup(base_model_path="stabilityai/stable-diffusion-xl-base-1.0",
         torch_dtype=torch.float16,
         add_watermarker=False,
     ).to(device)
-
-    # Register cross-attention hook for IP Adapter
     pipe.unet = register_cross_attention_hook(pipe.unet)
-    
-    # Initialize IP Adapter model
+
     ip_model = IPAdapterXL(pipe, image_encoder_path, ip_ckpt, device)
-    
-    # Load LoRA weights into the Stable Diffusion pipeline
-    pipe.load_lora_weights(
-        lora_model_path,
-        weight_name="PixelArt_v1_LoRA_XL.safetensors",  # Replace with your actual LoRA filename
-        adapter_name="pixelart"
-    )
 
-    # Load the VAE (Variational Autoencoder) for handling image decoding
-    vae = AutoencoderKL.from_pretrained(base_model_path, torch_dtype=torch.float16).to(device)
-    pipe.vae = vae
+    """
+    Get Depth Model Ready
+    """
 
-    # Set up the DPT Depth Model for depth map generation
     model = DPTDepthModel(
         path=model_depth_path,
         backbone=depth_backbone,
         non_negative=True,
         enable_attention_hooks=False,
-    ).to(device)
+    )
+
     model.eval()
 
-    return [ip_model, model, pipe]  # Return models including the pipeline with LoRA
+    return [ip_model, model]