Update processing/setup.py

processing/setup.py

@@ -1,71 +1,52 @@
 import huggingface_hub
 import torch
-from diffusers import ControlNetModel, StableDiffusionXLControlNetInpaintPipeline, DDIMScheduler, AutoencoderKL
+from diffusers import ControlNetModel, StableDiffusion3Pipeline
+
 from DPT.dpt.models import DPTDepthModel
-from ip_adapter import IPAdapter, IPAdapterXL
+from ip_adapter import IPAdapterXL
 from ip_adapter.utils import register_cross_attention_hook
 
-def setup(base_model_path="stabilityai/stable-diffusion-xl-base-1.0",
+
+def setup(base_model_path="stabilityai/stable-diffusion-3.5-medium",
           image_encoder_path="sdxl_models/image_encoder",
-          ip_ckpt="sdxl_models/ip-adapter_sdxl.bin",
-          controlnet_path="diffusers/controlnet-depth-sdxl-1.0",
+          ip_ckpt="sdxl_models/ip-adapter_3.5.bin",  # Update for 3.5
+          controlnet_path="diffusers/controlnet-depth-sd3.5",  # Updated path
           device="cuda",
           model_depth_path="DPT/weights/dpt_hybrid-midas-501f0c75.pt",
           depth_backbone="vitb_rn50_384"):
-    """Set up the processing module."""
+    """Set up the processing module for Stable Diffusion 3.5."""
     huggingface_hub.snapshot_download(
-        repo_id='h94/IP-Adapter',
-        allow_patterns=['models/**', 'sdxl_models/**'],
+        repo_id='stabilityai/stable-diffusion-3.5',
+        allow_patterns=[
+            'models/**',
+            'sd3.5_models/**',
+        ],
         local_dir='./',
         local_dir_use_symlinks=False,
     )
 
     torch.cuda.empty_cache()
 
-    # # Load scheduler
-    # noise_scheduler = DDIMScheduler(
-    #     num_train_timesteps=1000,
-    #     beta_start=0.00085,
-    #     beta_end=0.012,
-    #     beta_schedule="scaled_linear",
-    #     clip_sample=False,
-    #     set_alpha_to_one=False,
-    #     steps_offset=1,
-    # )
-
-    # Load VAE
-    vae_model_path = "stabilityai/sd-vae-ft-mse"
-    vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
-
-    # Load ControlNet model with depth conditioning
+    # Load Stable Diffusion 3.5 pipeline with ControlNet for depth
     controlnet = ControlNetModel.from_pretrained(controlnet_path, use_safetensors=True,
                                                  torch_dtype=torch.float16).to(device)
-    controlnet.conditioning_scale = 1.0  # Optional: Adjust as needed
-
-    # Load SDXL pipeline with additional components
-    pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(
+    pipe = StableDiffusion3Pipeline.from_pretrained(
         base_model_path,
         controlnet=controlnet,
         use_safetensors=True,
-        torch_dtype=torch.float16,
-        # scheduler=noise_scheduler,
-        vae=vae,
-        add_watermarker=False,
+        torch_dtype=torch.float16
     ).to(device)
-    
-    # Register cross-attention hook for IP Adapter
     pipe.unet = register_cross_attention_hook(pipe.unet)
 
-    # Load IP Adapter
     ip_model = IPAdapterXL(pipe, image_encoder_path, ip_ckpt, device)
 
-    # Initialize Depth Model
+    # Initialize DPT Depth Model
     model = DPTDepthModel(
         path=model_depth_path,
         backbone=depth_backbone,
         non_negative=True,
         enable_attention_hooks=False,
-    ).to(device)
+    )
     model.eval()
 
-    return [ip_model, model]
+    return [ip_model, model]