SDXL 77 token 이상 생성하기
SDXL로 이미지를 만들다 보면 프롬프트를 꽤 길게 썼는데도, 뒤에 적은 키워드가 결과에 거의 반영되지 않는 경우가 있다.
처음에는 모델이 프롬프트를 제대로 이해하지 못하는 것처럼 보이지만, 실제 원인은 조금 다르다.
핵심은 CLIP 텍스트 인코더가 한 번에 처리할 수 있는 토큰 수가 77개로 제한되어 있다는 점이다.
프롬프트가 이 길이를 넘으면 뒤쪽 내용이 잘리거나 약하게 반영될 수 있다.
이때 Compel 라이브러리를 사용하면 긴 프롬프트를 여러 조각으로 나눠 처리하는 방식으로 이 제한을 우회할 수 있다.
77 토큰 제한은 어디서 나오는가
positional embedding이 필요한 이유
Transformer 구조는 기본적으로 단어의 순서를 스스로 알지 못한다.
attention 연산은 토큰끼리의 관계를 계산하지만, 각 토큰이 몇 번째 위치에 있는지는 별도로 알려줘야 한다.
이 역할을 하는 것이 positional embedding이다.
쉽게 말하면, 각 위치마다 붙는 위치 정보라고 보면 된다.
0번째 자리에는 0번 위치 벡터, 1번째 자리에는 1번 위치 벡터가 더해지는 식이다.
텍스트가 토큰으로 바뀌고, 다시 임베딩 벡터로 변환되면 여기에 위치 정보가 더해진다.
이 과정을 통해 모델은 단어의 순서를 어느 정도 구분할 수 있게 된다.
CLIP은 77칸짜리 위치 테이블을 사용한다
문제는 CLIP이 학습될 때 positional embedding 테이블을 77칸 크기로 만들었다는 점이다.
즉, 0번부터 76번 위치까지만 학습된 값이 있다.
77번째 이후의 위치에 대해서는 모델이 학습한 위치 정보가 없다.
그래서 단순히 더 긴 프롬프트를 그대로 넣는다고 해결되지 않는다.
학습된 위치 정보가 없는 자리에 토큰을 추가해도, CLIP은 그 토큰을 제대로 해석하기 어렵다.
실제로 사용할 수 있는 건 75토큰
77칸을 모두 프롬프트 내용에 사용할 수 있는 것도 아니다.
양쪽 끝에는 특수 토큰이 들어간다.
맨 앞: <|startoftext|>
맨 뒤: <|endoftext|>따라서 실제 프롬프트 내용이 들어갈 수 있는 공간은 75토큰이다.
SDXL은 인코더가 두 개지만, 한도가 두 배가 되지는 않는다
SDXL은 CLIP-L과 OpenCLIP-G, 두 개의 텍스트 인코더를 사용한다.
그래서 처음에는 77개씩 두 개니까 총 154토큰까지 처리할 수 있을 것처럼 보일 수 있다.
하지만 실제 구조는 그렇지 않다.
두 인코더가 프롬프트를 절반씩 나눠 받는 것이 아니라, 같은 프롬프트를 각각 병렬로 처리한다.
따라서 토큰 한도가 합쳐지는 것이 아니라, 각 인코더가 따로 77토큰 제한을 가진다.
Compel은 이 제한을 어떻게 우회할까?
Compel의 방식은 생각보다 단순하다.
CLIP에 한 번에 77토큰까지만 넣을 수 있다면, 프롬프트를 여러 조각으로 나눠서 넣으면 된다.
그리고 CLIP을 통과한 결과 임베딩을 다시 이어 붙인다.
예를 들어 150토큰짜리 프롬프트가 있다고 가정해보자.
- 프롬프트를 75토큰 단위로 나눈다.
- 각 조각 앞뒤에 BOS와 EOS 토큰을 붙인다.
- 각 조각을 CLIP에 따로 넣는다.
- CLIP에서 나온 임베딩을 시퀀스 방향으로 이어 붙인다.
- 결과적으로 더 긴 프롬프트 정보를 담은 임베딩 시퀀스가 만들어진다.
형태로 보면 대략 이런 식이다.
[1, 77, 768] + [1, 77, 768]
→ [1, 154, 768]CLIP 입장에서는 매번 정상적인 77토큰 입력만 받는 셈이다.
하지만 CLIP을 통과한 뒤에는 각 조각의 임베딩을 이어 붙였기 때문에, 전체 프롬프트의 정보가 더 많이 남는다.
U-Net은 긴 임베딩을 받아도 괜찮을까?
여기서 한 가지 의문이 생길 수 있다.
CLIP은 77토큰 제한이 있는데, 그 뒤에 있는 U-Net은 더 긴 임베딩을 받아도 괜찮을까?
결론부터 말하면, 이 방식은 U-Net 쪽에서는 비교적 자연스럽게 처리된다.
77토큰 제한은 CLIP 내부의 positional embedding 때문에 생기는 문제다.
프롬프트가 CLIP을 통과해서 이미 임베딩 벡터가 되고 나면, 그 이후 단계에서는 CLIP의 위치 테이블 제한이 직접적으로 적용되지 않는다.
U-Net은 cross-attention을 통해 텍스트 임베딩을 참고한다.
이때 텍스트 시퀀스 길이가 77이든 154이든, 길이에 맞춰 attention 계산 크기가 달라질 뿐이다.
다시 말해 병목은 U-Net이 아니라 CLIP 쪽에 있다.
CLIP에 들어가기 전까지만 77토큰 제한을 맞춰주면, 그 이후에는 더 긴 임베딩 시퀀스를 넘기는 방식으로 처리할 수 있다.
Compel은 이 구조를 이용한다.
프롬프트를 CLIP에 넣기 전에는 여러 조각으로 나누고, CLIP을 통과한 뒤에는 나온 임베딩을 다시 이어 붙이는 방식이다.
코드로 확인하기
Compel 초기화
from compel import Compel, ReturnedEmbeddingsType
compel = Compel(
tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
requires_pooled=[False, True],
truncate_long_prompts=False
)여기서 가장 중요한 옵션은 truncate_long_prompts=False다.
이 값을 False로 둬야 75토큰을 넘는 프롬프트가 잘리지 않고, Compel이 내부적으로 프롬프트를 청크 단위로 나눠 처리한다.
tokenizer와 text_encoder를 리스트로 넘기는 이유는 SDXL이 두 개의 텍스트 인코더를 사용하기 때문이다.
두 인코더 모두 같은 프롬프트를 기준으로 임베딩을 만들어야 한다.
returned_embeddings_type에는 PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED를 사용했다.
마지막 레이어가 아니라 마지막에서 두 번째 레이어의 hidden state를 가져오는 설정이다.
requires_pooled=[False, True]는 pooled embedding을 어느 인코더에서 가져올지 정하는 옵션이다.
SDXL에서는 두 번째 인코더 쪽 pooled embedding이 필요하므로 두 번째 값만 True로 둔다.
positive와 negative 프롬프트 길이 맞추기
conditioning, pooled = compel(prompt)
negative_conditioning, negative_pooled = compel(negative_prompt)
conditioning, negative_conditioning = pad_to_same_length(
[conditioning, negative_conditioning]
)positive 프롬프트와 negative 프롬프트의 길이는 항상 같지 않다.
예를 들어 positive 프롬프트는 110토큰이고, negative 프롬프트는 30토큰일 수 있다.
이 경우 Compel을 거치면 positive 쪽은 154 길이, negative 쪽은 77 길이가 될 수 있다.
이 상태 그대로 classifier-free guidance에 넣으면 텐서 shape가 맞지 않아 에러가 날 수 있다.
그래서 짧은 쪽을 0으로 패딩해서 두 임베딩의 길이를 맞춰준다.
def pad_to_same_length(tensors):
max_len = max(t.shape[1] for t in tensors)
padded = []
for t in tensors:
pad_len = max_len - t.shape[1]
if pad_len > 0:
pad = torch.zeros(
(t.shape[0], pad_len, t.shape[2]),
device=t.device,
dtype=t.dtype
)
t = torch.cat([t, pad], dim=1)
padded.append(t)
return padded0 벡터로 패딩해도 결과에 큰 영향을 주지 않는 편이다.
cross-attention에서 0 벡터는 query와의 내적값이 의미 있게 잡히기 어렵기 때문이다.
파이프라인에 넣을 때 주의할 점
image = pipe(
prompt_embeds=conditioning,
pooled_prompt_embeds=pooled,
negative_prompt_embeds=negative_conditioning,
negative_pooled_prompt_embeds=negative_pooled,
...
).images[0]여기서 중요한 점은 prompt=에 문자열을 그대로 넘기면 안 된다는 것이다.
문자열 프롬프트를 그대로 넘기면 diffusers 내부에서 다시 토크나이저를 돌린다.
그러면 결국 기본 경로를 타면서 77토큰 제한에 걸릴 수 있다.
Compel로 긴 프롬프트 임베딩을 만들었다면, 파이프라인에는 prompt_embeds=로 직접 넘겨야 한다.
SDXL에서는 pooled_prompt_embeds도 함께 넘겨야 한다.
정리
SDXL에서 긴 프롬프트의 뒤쪽 내용이 잘 반영되지 않는 이유는 대부분 CLIP의 77토큰 제한과 관련이 있다.
Compel은 이 문제를 프롬프트 청크 분할 방식으로 우회한다.
프롬프트를 75토큰 단위로 나눠 CLIP에 따로 넣고, CLIP에서 나온 임베딩을 다시 이어 붙인다.
중요한 점은 문자열 프롬프트를 그대로 넘기지 않는 것이다.
긴 프롬프트를 제대로 사용하려면 truncate_long_prompts=False를 설정하고, 생성 단계에서는 prompt_embeds와 pooled_prompt_embeds를 직접 넘겨야 한다.
이렇게 하면 CLIP의 입력 제한은 지키면서도, 최종적으로는 더 긴 프롬프트 정보를 담은 임베딩을 U-Net에 전달할 수 있다.