Apple Watch에서 LCM Diffusion을 구동하여 256x256 이미지 생성까지 성공한 이야기

서론

Apple Watch 상에서 네트워크 연결 없이 이미지 생성을 구동하는 실험을 계속하고 있습니다.

지난번에는 Core ML을 사용하지 않고 순수 Swift의 작은 MLP로 이미지를 직접 생성하는 부분까지 정리했습니다.

이 기사는 그 후속편입니다.

이번에는 MLP 방식이 아니라, Apple Watch 상에서 LCM 계열의 작은 diffusion pipeline을 구동하여, 짧은 자유 입력 프롬프트(prompt)로부터 256x256 이미지를 생성하는 단계까지 도달한 이야기를 정리합니다.

결론부터 말씀드리면, SDXL과 같은 완전 자유 입력은 아니지만, 짧은 프롬프트라면 "이미지만 봐도 대략 무엇을 입력했는지 알 수 있을" 정도까지는 왔습니다.

구현 결과물

현재 watchOS 실기기용 baseline은 다음과 같은 구성입니다.

항목	내용
scheme	`WatchPipelineSmokeApp`
...
실기기 확인에 사용한 환경은 다음과 같습니다.

항목	내용
device	Apple Watch Series 11 (GPS)
...
파이프라인은 대략 다음과 같은 흐름입니다.

Apple Watch 실기기에서는 text encoder를 생성 파이프라인과 동시에 유지하지 않도록 하고 있습니다.

prompt를 tokenize한다
CLIP text encoder를 로드한다
[1, 77, 768]의 hidden_states를 만든다
text encoder를 해제한다
Core ML cache를 purge한다
LCM UNet chunk를 순서대로 로드/추론한다
decoder로 256x256 이미지로 되돌린다

이 "분리 실행"이 상당히 중요했습니다.

모델 결과물은 Hugging Face로 분리

.mlmodelc나 .mlpackage는 크기가 상당히 커집니다.

이번 LCM256 baseline만 해도, Watch app에 두는 compiled assets가 약 879MB, Mac 평가나 재빌드용 Core ML packages도 약 878MB에 달합니다.

따라서 GitHub repository 본체에는 넣지 않고, Hugging Face에 외부 artifact로 두기로 했습니다.

GitHub 측에는 코드, Xcode project, scheduler, tokenizer metadata, docs, 취득 스크립트만 둡니다.

clone한 사람은 다음 명령어로 Watch 실행에 필요한 모델을 취득할 수 있습니다.

python3 tools/fetch_watch_lcm256_assets.py

Mac 측에서 품질 평가도 하고 싶은 경우에는 .mlpackage도 취득합니다.

python3 tools/fetch_watch_lcm256_assets.py --packages

이렇게 함으로써 GitHub repo를 가볍게 유지하면서도, 다른 사람이 실기기/로컬 평가를 재현할 수 있는 형태로 만들 수 있습니다.

왜 MLP에서 diffusion으로 돌아갔는가

지난번의 MLP 방식은 Apple Watch에서 구동한다는 의미에서는 상당히 다루기 쉬운 방식이었습니다.

의존 모델 없음
구현이 작음
메모리 사용량이 적음
Xcode에서 즉시 빌드 가능

반면, 품질 면에서는 한계도 보였습니다.

학습된 어휘에서 벗어나면 약함
구도나 형상이 흐릿해지기 쉬움
이미지의 다양성이 낮음
짧은 자유 입력 prompt에 대응하기 어려움

그래서 MLP는 "경량 demo / 비교 baseline"으로 남겨두면서, 품질 측면의 본래 목표로서 LCM 계열의 diffusion pipeline을 시도하기로 했습니다.

우선 Core ML Stress로 한계를 확인

갑자기 생성 pipeline을 구성하기 전에, 별도의 scheme으로 Core ML의 load/predict/memory 한계를 측정했습니다.

사용한 격리 스킴(scheme)은 다음과 같습니다.

이 단계에서는 기존의 작은 Stable Diffusion 계열 Core ML 모델을 watchOS 실기기에 올려서, 단독 로드나 추론을 시도하고 있습니다.

메모리만 확인하는 Fine ladder에서는 16MB 단위로 272MB까지는 안정적이었으나, 288MB 부근에서 크래시가 발생했습니다. 이전의 aggressive ladder에서도 300MB 전후에서 떨어졌습니다.

즉, 대략적으로 말하자면 Apple Watch 상에서는 300MB 정도가 상당히 강력한 벽이 되고 있었습니다.

또한, MLComputeUnits.all를 사용했을 때는 UNet의 ANE 컴파일 실패나 메모리 급증 현상이 발생했습니다.

MILCompilerForANE error
_ANECompiler : ANECCompile() FAILED

따라서 이후의 Apple Watch 실기(device) 루트는 기본적으로 CPU-only로 진행하기로 했습니다.

이 스트레스 테스트(stress test)를 통해 알게 된 점은 크게 세 가지입니다.

파일 크기보다 로드(load) 시의 일시적 메모리와 Core ML 실행 계획(execution plan)이 더 위험하다
디코더(decoder) 단독으로는 상당히 가볍다
거대한 UNet을 한꺼번에 보유하는 것보다, 분할하여 일시적으로 로드하는 것이 현실적이다

LCM을 선택한 이유

일반적인 Stable Diffusion은 보통 20~30 step 정도를 전제로 하는 경우가 많아 Apple Watch에는 너무 무겁습니다.

반면 LCM은 적은 step으로 이미지를 생성할 수 있습니다. 이번 베이스라인(baseline)에서는 4 step을 사용했습니다.

물론 LCM을 선택했다고 해서 모든 문제가 해결되는 것은 아닙니다. UNet은 여전히 크고, 텍스트 인코더(text encoder)도 무겁습니다.

그래서 다음과 같이 분해했습니다.

텍스트 인코더(text encoder)는 일시적으로만 로드한다
UNet은 16개의 Core ML chunk로 분할한다
디코더(decoder)는 4bit로 양자화한다
모든 과정을 CPU-only로 구동한다
이미지는 256x256으로 고정한다
UI는 프롬프트(prompt) 입력과 생성 결과에만 집중한다

이 구성으로 Apple Watch 실기 상에서도 생성 단계까지 도달할 수 있었습니다.

128px, 192px, 256px

처음에는 64px나 128px부터 확인했습니다.

128px에서도 작동은 했지만, 외관상으로는 아직 상당히 작아서 무엇이 그려져 있는지 알 수 있는 것과 없는 것의 차이가 컸습니다.

그 후 192px를 시도하니 품질이 상당히 향상되었습니다.

나아가 256px도 시도해 본 결과, 실기 메모리 피크(peak)는 최대 약 140MB 정도로 수렴하였고, 품질도 한 단계 더 높아졌습니다.

생성 시간은 Apple Watch 실기에서 대략 1분 전후입니다.

최신 실기 로그에서는 cat mascot, blue bird, lion, ramen bowl, bicycle, ballerina 총 6회 실행 시, 합계 생성 시간은 53.8~62.8초, 평균 56.7초였습니다.

내역을 살펴보면, 텍스트 인코더(text encoder)의 로드와 추론(inference)은 대체로 2~~3초, 디코더(decoder)는 약 4초입니다. LCM 본체는 첫 번째 step에서 chunk 로드가 무겁고, 이후의 step은 8~~11초 정도로 안정화되는 형태였습니다.

256px로 설정함으로써 다음과 같은 카테고리들이 상당히 읽기 쉬워졌습니다.

동물
음식
자연 풍경
단일 오브젝트(object)
일부 로고/아이콘
일부 관계형 프롬프트(relationship prompt)

Mac 상에서 296장의 품질 평가

품질 평가는 Apple Watch 실기가 아닌 Mac 상에서 진행하고 있습니다.

그 이유는 모델 품질 비교만 목적이라면 Watch의 ANE나 실기 메모리는 필요하지 않으며, 동일한 스케줄러(scheduler), 토크나이저(tokenizer), 텍스트 인코더(text encoder), UNet, 디코더(decoder), 시드(seed) 규칙을 사용한다면 경향성을 파악하기에 충분하기 때문입니다.

평가용 스크립트는 다음과 같습니다.

실행 명령어는 다음과 같은 형태입니다.

.venv/bin/python tools/watch_lcm256_quality_eval.py \
--out-dir reports/watch_lcm256_quality/full_lcm256_g6

평가는 74 prompts x 4 seeds = 296 images로 진행했습니다.

결과는 다음과 같습니다.

항목	결과
생성 매수	296
...

상세한 평가 메모는 여기에 두었습니다.

생성 예시

다음은 Mac 평가에서 생성한 이미지의 발췌본입니다.

Apple Watch 실기 스크린샷이 아니라, 동일한 Core ML package를 Mac에서 CPU-only로 실행한 품질 평가 이미지입니다. 실기 측에서는 최종적인 메모리, 시간, 발열, UI를 확인하고 있습니다.

개인적으로는 이 정도 수준까지 나온다면 "Apple Watch 상의 로컬 이미지 생성 데모(demo)"로서 상당히 흥미로운 단계까지 왔다고 생각합니다.

특히 다음과 같은 프롬프트(prompt)는 상당히 잘 식별됩니다.

snowy mountain

strawberry cake

blue bird

dragon

astronaut riding a horse

cat in a spaceship

Apple logo

steam train

실기 스크린샷에서는 예를 들어 다음과 같은 출력이 나왔습니다.

Mac 평가 이미지보다 UI가 포함된 모습이지만, lion, cat mascot, blue bird, ramen bowl 정도는 Watch 화면상에서도 상당히 읽기 쉬운 결과로 나타나고 있습니다.

seed reroll이 중요함

Watch UI에서는 동일한 prompt로 재생성할 경우 seed를 reroll하도록 구현했습니다.

이것은 상당히 중요했습니다.

동일한 prompt라도 seed에 따라 결과가 크게 달라집니다.

예를 들어 cat in a spaceship은 seed에 따라 '무언가 푸르게 빛나는 고양이 같은 것'부터 '상당히 고양이와 우주선 같은 것'까지 변합니다.

bicycle과 같은 가는 구조는 취약하지만, 그럼에도 seed에 따라 알아볼 수 있는 형태가 됩니다.

이 모델에서는 한 번에 완벽한 이미지를 내놓는 것보다, 짧은 prompt를 넣고 여러 번 reroll하는 경험이 더 자연스럽습니다.

취약한 부분

물론 아직 무엇이든 생성할 수 있는 것은 아닙니다.

취약한 사례도 있습니다.

취약한 부분은 대체로 다음과 같습니다.

가는 구조
- bicycle
- 손가락, 팔, 다리
- 악기의 현이나 세부 묘사
pose 의존적인 인물
- ballerina
- superhero
추상적인 style 용어
- neon spiral
- origami crane
엄격한 관계 표현
- holding
- playing
- wearing
정확한 글자나 로고

다만, 수치상 very_soft로 판정된 이미지가 육안으로는 나쁘지 않은 경우도 있습니다.

예를 들어 green apple이나 moonlit lake는 에지(edge)가 적기 때문에 자동 평가에서는 낮게 나오지만, 이미지로서는 그렇게 나쁘지 않습니다.

따라서 현시점의 평가에서는 자동 스코어는 어디까지나 이상치 탐지에 사용하고, 최종 판단은 contact sheet의 육안 확인으로 하고 있습니다.

296장 평가를 통해 보인 경향

장르별로는 대체로 다음과 같은 인상입니다.

장르	특징
nature	가장 안정적. mountain, ocean, beach, waterfall 등이 강함
...

제로 플래그(zero flag)에서 안정적이었던 prompt에는 다음과 같은 것들이 있습니다.

dolphin

tiger

dragon

phoenix bird

hamburger

ramen bowl

strawberry cake

sushi plate

cat badge

snowy mountain

ocean wave

waterfall

leather boot

astronaut riding a horse

steam train

yellow bus

반면, 4개의 seed 모두에서 취약했던 prompt도 있었습니다.

ballerina

green apple

moonlit lake

neon spiral

이 부분은 향후 학습 데이터나 prompt 전개 개선 대상이 될 것 같습니다.

실기 측에서 주의한 점

Apple Watch 실기에서는 모델이 동작하는지 여부뿐만 아니라, 다음 사항들을 상당히 신경 썼습니다.

text encoder를 유지하지 않음

CLIP text encoder는 크기가 크기 때문에, 생성 중에 UNet과 함께 보유하는 것을 피했습니다.

prompt embedding을 만든 후 즉시 해제합니다.

UNet을 분할함

UNet을 한꺼번에 보유하는 것이 아니라, 16개의 chunk로 나누어 순차적으로 통과시킵니다.

이를 통해 로드(load) 시의 메모리 피크(peak)를 억제할 수 있습니다.

CPU-only로 설정

ANE나 GPU를 사용하면 빨라질 가능성은 있지만, watchOS 실기에서는 Core ML 컴파일(compile)이나 메모리 동작이 불안정해지기 쉽습니다.

현재는 CPU-only 방식이 더 읽기 쉽고 제어하기 용이합니다.

Xcode를 통한 실기 설치도 무거움

이 구성은 생성 중의 메모리 사용량뿐만 아니라, 앱 설치 자체도 무겁습니다.

한 번 Watch 상의 앱을 삭제한 후 다시 빌드해 보니, 실기기 상의 앱 크기는 1.57GB였습니다. 확인 당시 Watch의 여유 공간은 약 44GB입니다.

그렇기 때문에 Xcode에서 실기기로 빌드/설치하는 것만으로도 10분 이상 걸립니다. 환경에 따라서는 한 번에 설치 완료까지 진행되지 않을 가능성도 있습니다.

품질 평가를 대량으로 돌릴 경우에는 Mac 상의 평가 스크립트를 사용하고, Watch 실기기에서는 "실제로 로드할 수 있는지", "메모리에 들어가는지", "UI로서 기다릴 만한지", "착용 시 뜨겁게 느껴지지 않는지"를 확인하는 방식으로 역할을 나누는 것이 다루기 쉽습니다.

또한, Xcode에서 실행하여 콘솔 로그 (console log)를 보면서 구동하는 경우는 괜찮지만, Watch 상에서 앱 단독으로 실행하면 생성 중에 화면 표시가 꺼지면서 watchOS 측에서 앱이 일시 정지(suspend) 또는 종료(kill)되는 경우가 있습니다.

이는 아마도 Apple Watch 측의 앱 라이프사이클 (app lifecycle) 영향일 것입니다. WKExtendedRuntimeSession 등으로 개선할 수 있는 가능성은 있지만, 용도 제약도 있을 것으로 보여 이 기사 시점에서는 미대응 상태입니다.

UI를 작게 만들기

처음에는 디버깅용으로 preset, seed, guidance, preview mode 등을 UI에 표시했습니다.

최종적으로는 Watch 상에서 다음 항목들로만 집중했습니다.

prompt 입력란
Generate / Reroll Seed 버튼
생성 결과

상세 내용은 Xcode 콘솔 (console)에 로그로 출력합니다.

향후 하고 싶은 것

여기서 품질을 크게 높이려면 단순히 해상도를 높이는 것만으로는 어려워 보입니다.

256px화로 상당히 개선되었지만, 그 이상은 시간이나 활성화 메모리 (activation memory) 대비 효율이 낮을 것 같습니다.

다음 개선 후보는 다음과 같습니다.

LCM256을 위한 추가 학습/증류 (distillation)
Watch용 디코더 (decoder) 개선
짧은 prompt에 특화된 경량 텍스트 인코더 (text encoder)
prompt expansion 개선
seed 후보의 다중 생성/선택 UI
배터리 소모 (battery drain), 발열 (thermal), 연속 생성 시의 실기기 평가

특히, bicycle, ballerina, holding과 같은 약점은 해상도뿐만 아니라 학습 데이터 측면의 개선이 필요해 보입니다.

요약

Apple Watch 상에서 LCM 계열 diffusion을 구동하여, 짧은 자유 입력 prompt로부터 256x256 이미지를 생성할 수 있는 단계까지 왔습니다.

현시점의 주요 포인트는 다음과 같습니다.

Apple Watch 실기기에서 LCM256 6bit 파이프라인 (pipeline)이 작동함
CLIP text encoder도 분리 실행이라면 사용 가능함
UNet은 16 chunk로 나누어 스트리밍 (streaming)함
CPU-only Core ML로 안정성을 우선함
실기기 메모리 피크는 약 140MB 정도로 억제됨
실기기 생성 시간은 256px 기준 대략 1분 내외였음
Mac 품질 평가에서는 296장을 생성하여 실패 0건을 기록함
짧은 prompt라면 이미지 자체만으로도 의도를 알 수 있는 경우가 늘어남

아직 SDXL처럼 자유로운 입력은 아닙니다.

하지만 Apple Watch 단독으로 이 정도 이미지가 나온다면, 기술 데모 (technical demo)로서 상당히 흥미로운 단계에 진입했다고 생각합니다.

리포지토리는 여기 있습니다.