depth lora causing transformer layers mismatch #11109
Description
Describe the bug
When the depth lora is injected into the base dev model, peft changes the transformer's linear layer self.x_embedder from (3072, 64) to (3072, 128). These peft changes persist even after unloading the depth lora , causing subsequent requests to throw a layer mismatch error.
Reproduction
import torch
from diffusers import FluxControlPipeline, FluxPipeline
from diffusers.utils import load_image
pipeline = FluxControlPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, cache_dir = "/home/user/train_model_shared/huggingface").to("cuda")
print(f"transformer in channels: " , pipeline.transformer.config.in_channels)
pipeline.load_lora_weights("black-forest-labs/FLUX.1-Depth-dev-lora", adapter_name="depth", cache_dir = "/home/user/train_model_shared/huggingface")
print(f"transformer in channels: " , pipeline.transformer.config.in_channels)
pipeline.set_adapters("depth", 0.85)
prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."
control_image = load_image("depth.png")
image = pipeline(
prompt=prompt,
control_image=control_image,
height=1024,
width=1024,
num_inference_steps=30,
guidance_scale=10.0,
generator=torch.Generator().manual_seed(42),
).images[0]
image.save("output1.png")
pipeline.unload_lora_weights()
print(f"transformer in channels: " , pipeline.transformer.config.in_channels)
# Move the pipeline to CUDA and convert to bfloat16
pipeline = FluxPipeline(**pipeline.components).to("cuda")
pipeline.to(dtype=torch.bfloat16)
image = pipeline(
prompt=prompt,
height=1024,
width=1024,
num_inference_steps=30,
guidance_scale=10.0,
generator=torch.Generator().manual_seed(42),
).images[0]
image.save("output2.png")
Logs
This is the error that i am getting on this script, but i think the main issue is transformers layer mismatch.
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/user/runware/shehzad/temp/sd-base-api/diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py", line 446, in step
prev_sample = sample + dt * model_output
~~~~~~~^~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (128) must match the size of tensor b (64) at non-singleton dimension 2System Info
Diffusers: 0.33.0.dev0
python: 3.12