Published on

Stable Diffusion LoRA 깨짐? UNet 키 불일치 해결

Authors

Stable Diffusion에서 LoRA를 적용했는데 갑자기 로드 에러가 나거나, 에러 없이 로드는 되는데 결과물이 심하게 깨지는 경우가 있습니다. 로그를 보면 대개 Missing key(s) in state_dict / Unexpected key(s) 같은 메시지와 함께 UNet 쪽 키가 맞지 않는다는 신호가 뜹니다.

이 글에서는 UNet 키 불일치가 왜 생기는지, SD 1.5 vs SDXL vs 파생 모델에서 어떤 지점이 달라지는지, 그리고 diffusers / AUTOMATIC1111 / ComfyUI에서 실무적으로 통하는 해결 루틴을 정리합니다.


증상: “LoRA가 깨졌다”가 아니라 “키가 안 맞는다”

대표 증상은 아래 중 하나로 나타납니다.

  • LoRA 로드 시 즉시 실패
    • 예: RuntimeError: Error(s) in loading state_dict for UNet2DConditionModel: Missing key(s)... Unexpected key(s)...
  • 로드는 되는데 출력이 망가짐
    • 얼굴이 뭉개지거나, 특정 스타일이 과하게 증폭되거나, 노이즈가 남는 현상
  • 같은 LoRA가 어떤 체크포인트에서는 되고, 다른 체크포인트에서는 실패

핵심은 LoRA 파일 내부의 타깃 모듈 이름(키 경로)현재 사용 중인 UNet 아키텍처의 모듈 이름(키 경로) 이 다르다는 점입니다.


원인 1: SD 1.5용 LoRA를 SDXL에 억지로 로드

가장 흔한 케이스입니다.

  • SD 1.5 계열 UNet은 attention 블록 구성과 채널 수가 SDXL과 다릅니다.
  • LoRA는 보통 UNet의 특정 레이어(예: attention의 to_q, to_k, to_v, to_out)에 저랭크 가중치를 덧씌우는 방식인데, 대상 레이어가 다르면 키가 맞을 수가 없습니다.

체크포인트가 SDXL인지 SD 1.5인지 헷갈린다면, 파일명만 보지 말고 메타데이터 또는 로더 로그로 확인하세요.


원인 2: LoRA 포맷 차이 (kohya-ss / diffusers / A1111)

LoRA는 “개념”은 같아도 저장 포맷과 키 네이밍이 다를 수 있습니다.

  • safetensors 내부 키가 A1111 스타일인 경우
  • diffusers용으로 변환된 LoRA
  • LyCORIS(LoCon/LoHa 등)처럼 확장 기법

특히 A1111은 UNet 키를 lora_unet_... 같은 접두로 평탄화해서 저장하는 경우가 있고, diffusers는 모듈 경로가 더 파이썬 객체 구조에 가깝습니다.


원인 3: 베이스 모델이 “파생 UNet” (inpainting, depth, refiner 등)

  • SD 1.5 inpainting 모델은 입력 채널이 달라 UNet 구조가 일부 다릅니다.
  • SDXL은 base와 refiner가 다르고, refiner 전용 LoRA도 존재합니다.

즉 “SD 1.5용 LoRA”라고 해도 일반 1.5에 맞춘 것인지, inpainting에 맞춘 것인지에 따라 키/shape가 달라질 수 있습니다.


진단 1: LoRA 내부 키를 직접 열어보기

가장 빠른 방법은 LoRA 파일에 어떤 키가 들어있는지 확인하는 겁니다. 아래 코드는 safetensors LoRA의 키 목록과 텐서 shape를 일부 출력합니다.

from safetensors.torch import load_file

path = "./my_lora.safetensors"
st = load_file(path)

print("num_keys:", len(st))
for i, (k, v) in enumerate(st.items()):
    if i > 30:
        break
    print(k, tuple(v.shape))

여기서 확인 포인트:

  • 키에 lora_unet_ 같은 접두가 있는지
  • lora_te_ 또는 text_encoder 관련 키가 있는지
  • unet 외에 text_encoder_2 같은 SDXL 흔적이 있는지

SDXL LoRA는 보통 텍스트 인코더 쪽도 text_encodertext_encoder_2로 분리되어 들어갈 수 있습니다.


진단 2: 현재 UNet의 state_dict 키와 비교

diffusers를 쓰고 있다면 UNet의 키를 뽑아 LoRA 키와 비교하면 “왜 불일치인지”가 명확해집니다.

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

unet_keys = set(pipe.unet.state_dict().keys())
print("unet_keys:", len(unet_keys))

# 일부 샘플 출력
for i, k in enumerate(sorted(unet_keys)):
    if i > 30:
        break
    print(k)

LoRA 키와 UNet 키의 경로 규칙이 아예 다르면, 단순히 “누락”이 아니라 네이밍 체계가 다르다는 뜻입니다.


해결 1: 가장 안전한 원칙 — “LoRA와 베이스 모델 계열을 맞춰라”

정답처럼 들리지만, 실제로 이것만 지켜도 문제의 80%가 사라집니다.

  • SD 1.5 LoRA는 SD 1.5 계열 체크포인트에서만 사용
  • SDXL LoRA는 SDXL base 또는 refiner 중 학습 대상에 맞춰 사용
  • inpainting 전용이면 inpainting 베이스에서 사용

만약 “이 LoRA가 어느 계열인지” 불명확하다면, Civitai 같은 배포 페이지의 base model 표기를 확인하거나, 위에서 설명한 키 검사로 SDXL 흔적(text_encoder_2)을 찾는 식으로 판별하세요.


해결 2: diffusers에서 LoRA 로드 시 adapter_name 분리와 스케일 관리

로드는 되는데 결과가 깨지는 경우는 스케일 과다 또는 여러 LoRA 중첩이 원인일 때가 많습니다.

diffusers 예시:

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

pipe.load_lora_weights("./lora_dir", weight_name="my_lora.safetensors", adapter_name="style")
pipe.set_adapters(["style"], adapter_weights=[0.7])

image = pipe("portrait photo, 85mm, soft light").images[0]
image.save("out.png")

체크 포인트:

  • 스케일을 1.0부터 박는 습관을 버리고, 0.4~0.8 구간에서 먼저 탐색
  • 여러 LoRA를 섞으면 각 스케일 합이 커져서 “깨짐”처럼 보일 수 있음

해결 3: “키 불일치”가 네이밍 문제라면 변환(매핑)로 해결

어떤 경우엔 LoRA 자체는 같은 계열인데, 저장 도구가 달라 키 접두/구분자가 달라서 로더가 못 알아듣는 경우가 있습니다. 이때는 변환 스크립트(키 리라이트)가 필요합니다.

아래는 키 문자열 치환을 통해 새 safetensors를 만드는 최소 예시입니다. (실제 매핑 규칙은 케이스마다 다르므로, 먼저 키 패턴을 비교하고 적용해야 합니다.)

from safetensors.torch import load_file, save_file

src = "./broken_lora.safetensors"
dst = "./fixed_lora.safetensors"

st = load_file(src)
new = {}

for k, v in st.items():
    nk = k
    # 예시: 특정 접두 제거
    if nk.startswith("lora_unet_"):
        nk = nk.replace("lora_unet_", "unet.")

    # 예시: 구분자 치환
    nk = nk.replace("_", ".")

    new[nk] = v

save_file(new, dst)
print("saved:", dst, "keys:", len(new))

주의:

  • 무작정 _.로 바꾸는 식의 변환은 대부분 실패합니다.
  • 반드시 “원래 UNet 키”와 “LoRA 키”를 나란히 놓고, 규칙적으로 매핑 가능한지 확인해야 합니다.
  • shape mismatch가 함께 뜨면 네이밍 문제가 아니라 아키텍처 자체가 다르다는 뜻일 가능성이 큽니다.

해결 4: A1111에서 UNet 키 불일치가 날 때 점검 체크리스트

A1111(또는 WebUI 계열)에서 흔히 놓치는 지점들입니다.

  1. LoRA가 LyCORIS 계열인지 확인
    • LyCORIS는 전용 확장 또는 로더가 필요할 수 있습니다.
  2. LoRA 파일이 손상되었는지 확인
    • 다운로드 중 깨지면 키 로드가 비정상적으로 보일 수 있습니다.
  3. 베이스 체크포인트가 “merge 모델”인지 확인
    • merge 과정에서 UNet이 바뀌거나, 예상과 다른 구조가 들어갈 수 있습니다.
  4. 콘솔 로그에서 실제 Missing/Unexpected 키를 복사해 패턴 분석
    • 어떤 블록(down_blocks, mid_block, up_blocks)에서 주로 터지는지에 따라 원인이 갈립니다.

이 과정은 “원인 찾기”가 핵심이라, 한 번에 해결하려고 하기보다 로그를 기반으로 분류하는 게 빠릅니다. 이런 접근은 캐시/빌드 문제를 진단할 때도 동일하게 통합니다. 예를 들어 CI에서 꼬인 상태를 강제로 초기화하는 패턴은 GitHub Actions 캐시 충돌 시 빌드 완전 초기화 전략에서 다룬 방식과 유사합니다.


해결 5: ComfyUI에서 “모델은 로드되는데 결과가 이상”할 때

ComfyUI는 노드 그래프가 유연한 대신, 아래 실수가 잦습니다.

  • SDXL base에 SDXL refiner용 LoRA를 연결
  • Clip 노드가 SDXL용(ClipG/ClipL 분리)인지, SD 1.5용인지 혼동
  • LoRA 적용 노드가 UNet과 CLIP에 동시에 적용되는데, 한쪽만 맞고 다른 쪽이 불일치

대응:

  • LoRA를 UNet에만 적용해보고 정상 여부 확인
  • CLIP 적용을 끄거나 분리해 어떤 쪽에서 깨지는지 확인
  • 스케일을 낮춰 “깨짐”이 과적용인지 확인

실전 루틴: 10분 안에 결론 내는 순서

  1. 베이스 모델 계열 확인: SD 1.5 vs SDXL vs inpainting
  2. LoRA 키 검사: safetensors 키에 SDXL 흔적(text_encoder_2)이 있는지
  3. 로더 로그에서 Missing key(s) / Unexpected key(s) 패턴 확인
  4. 같은 계열의 “표준 체크포인트”에 LoRA를 붙여 재현
    • SD 1.5면 runwayml/stable-diffusion-v1-5 같은 기준점
  5. 기준점에서 실패하면 LoRA 자체 문제(포맷/손상/다른 계열)
  6. 기준점에서 성공하면 현재 체크포인트 쪽 문제(merge, 파생 UNet, refiner 혼동)

이 루틴은 장애 분석에서 “기준 환경으로 재현”하는 방식과 같습니다. 예를 들어 서버/런타임 문제도 기준 환경을 잡아 원인을 좁히는 것이 중요하고, 그런 디버깅 사고방식은 Azure VM 부팅 실패? Boot Diagnostics로 5분 진단 같은 글에서 다루는 접근과도 닮아 있습니다.


자주 묻는 질문

LoRA 로드는 되는데 결과가 “노이즈처럼” 깨져요

  • 스케일이 너무 큰 경우가 많습니다. 0.6 전후로 낮춰보세요.
  • 서로 다른 목적의 LoRA를 여러 개 섞으면 과적용이 누적됩니다.
  • SDXL에서 base와 refiner 파이프라인을 섞어 쓰는 경우, refiner 단계에 맞지 않는 LoRA를 적용하면 결과가 흔들릴 수 있습니다.

Missing key(s)만 뜨고 Unexpected key(s)는 없어요

  • 보통 LoRA가 기대하는 레이어가 현재 UNet에 없다는 뜻입니다.
  • 체크포인트가 inpainting/파생 구조일 가능성이 있습니다.

변환 스크립트로 다 해결할 수 있나요

  • 네이밍 차이만 있으면 가능하지만, shape mismatch는 변환으로 해결이 어렵습니다.
  • shape mismatch는 사실상 “다른 모델 계열”이거나 “다른 UNet 구조”일 가능성이 큽니다.

마무리

LoRA가 “깨졌다”는 표현은 현상일 뿐이고, 실제로는 UNet 키 네이밍/구조 불일치라는 구체적인 원인이 있습니다.

  • 먼저 계열(SD 1.5, SDXL, inpainting, refiner)을 맞추고
  • safetensors 키를 열어 어떤 모듈을 겨냥하는지 확인한 다음
  • 기준 체크포인트에서 재현 테스트로 LoRA 문제인지 체크포인트 문제인지 분리

이 순서로 접근하면, 감으로 이것저것 바꾸는 것보다 훨씬 빠르게 결론에 도달합니다. 만약 키 불일치 로그를 정리하다가 “긴 작업을 추적해 병목을 찾아내는” 방식이 필요하다면, 성능 분석 관점에서는 Chrome INP 점수 급락? Long Task 추적·해결처럼 원인을 계측으로 좁히는 접근도 큰 도움이 됩니다.