파이톨치

[BoostCamp AI Tech] Diffusion 모델 개념부터 코드까지 살펴보기 본문

카테고리 없음

[BoostCamp AI Tech] Diffusion 모델 개념부터 코드까지 살펴보기

파이톨치 2024. 9. 5. 22:51
728x90

Diffusion 모델

Diffusion 모델은 수식이 좀 나온다. 

사실, 이론은 재미 없다. 하지만 알아야 한다. 하지만 이론부터 배우기엔 내 뇌는 참을성이 없다. 그러니 결과부터 보자. 

https://www.prpt.ai/prompt/imageDetail/520

대박이다. 진짜. 그림을 좋아하는 입장에서 저런 고퀄리티의 그림을 빠른 시간에 뽑아낼 수 있는 기술은 놀랍다. 

단순히 텍스트만 넣어주면 고퀄리티의 그림이 나오는 것이다. 사실 사용만 하는 입장에서는 몰라도 되는 것이다. 

하지만, 원리를 알면 더 재밌지 않을까? (아니 사실 원리 노잼이다.) 

 

Denoising Diffusion Probabilistic Models

Diffusion은 확산을 의미한다. 여기서 확산은 노이즈가 추가되는 것을 의미하는 모양이다.

이미지가 노이즈에서 시작해서 원래 이미지로 복구된다. 그 과정에서 한 스텝, 한 스텝에서 이미지를 복원하는 확률 분포를 학습하게 된다. 

(이게 너무 추상적이라 잘 이해는 안가지만...) 

이것은 픽셀 단위로 일어나는 계산이라 엄청 느리다고 한다. 때문에 다음에 Latent 벡터를 사용하는 방법이 나온다. 

 

Latent Diffusion(a.k.a. Stable Diffusion)

어렵다.  이미지 x가 인코더에 들어가면 z 벡터가 된다. 이미지 전체를 쓰는 것이 아니라, z 벡터를 쓰겠다는 컨셉이다. 

이미지를 압축한다고 생각하면 된다. stable diffusion의 장점은 conditioning이다. 저기에 text 등의 정보가 들어가는 것이다. 

굉장히 효과적이다!! 이제 text를 통해서 새로운 이미지를 생성할 수 있는 세상이 온 것이다. 위의 그림처럼 말이다!

Loss 는 사실 보면 MSE다. 

 Stable Diffusion 코드 뜯어보기

class StableDiffusion(nn.Module):
    def __init__(self, device, fp16, vram_O, sd_version='2.1', hf_key=None, t_range=[0.02, 0.98]):
        #############################################################
        # TODO: init 과정에서 반드시 들어가야 하는 super 관련 코드를 작성해주세요.

        #############################################################

        self.device = device  # device는 gpu에서 동작합니다
        self.sd_version = sd_version  # stable diffusion 2.1 version을 활용합니다

        print(f'[INFO] loading stable diffusion...')

        model_key = "stabilityai/stable-diffusion-2-1-base"

        self.precision_t = torch.float16 if fp16 else torch.float32

        # Stable diffusion model 선언
        pipe = StableDiffusionPipeline.from_pretrained(model_key, torch_dtype=self.precision_t)
        pipe.to(device)

        # Stable diffusion 모델의 구성요소인 VAE(variational autoencoder), UNet 구조를 선언해줍니다
        self.vae = pipe.vae
        self.tokenizer = pipe.tokenizer
        self.text_encoder = pipe.text_encoder
        self.unet = pipe.unet


        # DDPM의 발전된 버전인 DDIM scheduler를 활용해줍니다
        # DDIM scheduler는 해당 논문에서 확인 가능합니다: https://arxiv.org/abs/2010.02502
        self.scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler", torch_dtype=self.precision_t)

        del pipe

        # self.num_train_timesteps: 1000 으로, DDPM에서 활용하는 timestep 입니다
        # t_range는 [0.02, 0.98]으로 너무 작거나 큰 timestep이 샘플링되는 것을 막기 위해서 범위를 이와같이 제한하여 사용합니다
        self.num_train_timesteps = self.scheduler.config.num_train_timesteps

        #############################################################
        # TODO: t_range를 활용하여 int 값인 min_step과 max_step을 정의해주세요
        self.min_step = None  # None 지우고 코드 작성
        self.max_step = None  # None 지우고 코드 작성

        #############################################################

        self.alphas = self.scheduler.alphas_cumprod.to(self.device) # for convenience

        print(f'[INFO] loaded stable diffusion!')

맨 위에 들어갈 TODO는 알지? super().__init__()이다. 저기있는 model_key를 통해서 모델을 가져오는 모양이다. 

그 아래에서 pipe로써 pretrained된 디퓨전 모델을 가져온다. 당연히 pipe.to(device)를 해줘야 한다. 

저 스케쥴러는 2차원 디퓨전에서 사용하는 스케쥴러인 모양이다. pipe를 내장 변수에 다 넣어주면 지워준다. 

 

vae의 역할은 뭐지? chatgpt한테 물어보자. 아! 위 그림에 있는 인코더다! 

Stable Diffusion에서 VAE(Variational Autoencoder)는 이미지 데이터를 압축하고 복원하는 역할을 합니다. 구체적으로는, Stable Diffusion 모델에서 이미지와 잠재 공간(Latent Space) 간의 변환을 돕습니다. VAE는 이미지 데이터를 잠재 공간으로 변환하고, 다시 복원하는 과정을 통해 이미지 생성 과정에 중요한 역할을 합니다.
VAE의 역할
잠재 공간(Latent Space)로의 인코딩VAE의 인코더(encoder) 부분은 이미지 데이터를 압축하여 더 작은 차원의 잠재 공간으로 변환합니다.이 잠재 공간은 저차원의 벡터로, 이미지의 주요 특징을 압축해서 표현합니다. Stable Diffusion에서 이 잠재 공간은 모델이 처리하기에 더 효율적이며, 이미지 생성과 관련된 중요한 정보만을 유지하는 역할을 합니다.노이즈 제거 및 생성 과정에 활용Stable Diffusion 모델은 이 잠재 공간에서의 노이즈 제거 과정을 반복적으로 수행하여 최종 이미지를 생성합니다.노이즈가 점차 제거된 잠재 벡터는 VAE의 **디코더(decoder)**를 통해 다시 고해상도의 이미지로 복원됩니다. 즉, 디코더는 잠재 공간에서의 벡터를 이미지로 다시 변환하는 역할을 합니다.

 

즉, 이미지가 들어가면 sefl.vae(x) 이런 식으로 넘어가야 하는 것이다!  

아래 보니 이런 함수가 있었다. 그렇다면 이 함수는 어디서 써야 할까? 

    # 이미지를 VAE의 latent로 인코딩 해줍니다
    def encode_imgs(self, imgs):
        # imgs: [B, 3, H, W]

        imgs = 2 * imgs - 1

        posterior = self.vae.encode(imgs).latent_dist
        latents = posterior.sample() * self.vae.config.scaling_factor

        return latents

 

그것이 내가 채워야할 코드였다. 

        if as_latent:
            # 예측된 pred_rgb를 VAE를 통해 latent로 보내지 않고, pred_rgb 자체를 latent로 취급해서 바로 사용합니다
            # Hint: F.interpolate 함수 활용

            #############################################################
            # TODO: as_latent가 True일 경우에는, encoder를 통해서 dimension을 낮추는 것이 아니라,
            # pred_rgb 이미지 자체를 latent로 취급하고 image를 interpolation을 통해 사이즈를 줄여줍니다.
            # 예측된 pred_rgb를 입력으로 받아, (64, 64)로 사이즈를 조절하는 코드를 작성하세요.
            # latents의 범위는 (0,1)에서 (-1, 1)이 되도록 조정해주세요.

            latents = None  # None 지우고 코드 작성
            #############################################################
        else:
            # VAE encoder를 태우기 위해서는 pred_rgb resolution을 (512, 512)로 맞춰줘야 하기에 사이즈 변경을 위한 interpolation을 진행합니다
            # Hint: F.interpolate 함수 활용

            #############################################################
            # TODO: as_latent가 False일 경우에는, encoder를 통해서 dimension을 낮춰줍니다.
            # 예측된 pred_rgb를 입력으로 받아, (512, 512)로 사이즈를 가지는 pred_rgb_512를 만드는 코드를 작성하세요.
            # (Stable Diffusion encoder의 input dimension이 (512, 512) 입니다.)
            # pred_rgb_512를 클래스 메소드로 정의된 encode_imgs를 활용하여 latents로 인코딩하는 코드를 작성하세요.

            pred_rgb_512 = None  # None 지우고 코드 작성
            latents = None  # None 지우고 코드 작성

            #############################################################

as_latent는 이미 latent로 사용되는 모양이다. 그렇다면 latens = as pred_rgb가 될 것이다. 

latens는 (64, 64)가 되어야 하는 모양이다. 사이즈가 맞지 않을 수 있다. 또 무리해서 늘리거나 줄이면 안되기 때문에 보간을 해야한다. 

채운다는 의미이다. 

"interpolation" 코드
는 두 지점 사이의 값을 계산하는 데 주로 사용됩니다. 특히, 두 이미지, 벡터, 혹은 스칼라 값 사이를 선형 또는 비선형 방식으로 보간(interpolation)할 때 유용합니다. 이를 통해 중간 상태의 값이나 이미지를 생성할 수 있습니다.

이미지 크기를 변경하는 다양한 보간(interpolation) 기법을 통해 손실을 최소화하면서 이미지를 늘리거나 줄이는 방법이 있습니다. 일반적으로 사용되는 보간 기법은 다음과 같습니다:

최근접 이웃 보간(Nearest Neighbor Interpolation): 가장 단순한 방식으로, 가장 가까운 픽셀의 값을 사용해 새로운 픽셀을 결정합니다. 빠르지만 품질이 낮습니다.
양선형 보간(Bilinear Interpolation): 각 픽셀의 값을 주변 4개의 픽셀 값을 사용해 선형적으로 보간하는 방식입니다.
양입방 보간(Bicubic Interpolation): 주변 16개의 픽셀을 이용해 더 부드러운 이미지를 생성합니다. 고품질의 보간을 제공합니다.
           latents = pred_rgb  # None 지우고 코드 작성
            latents = F.interpolate(latents, size=(64, 64), mode='bilinear', align_corners=False)
            latents = latents - 0.5
            latents = latents * 2

이렇게 한다. 

pred_rgb가 아니라면 오히려 간단하다. 

            pred_rgb_512 = F.interpolate(pred_rgb, size=(512, 512), mode='bilinear', align_corners=False)  # None 지우고 코드 작성
            latents = self.encode_imgs(pred_rgb_512)  # None 지우고 코드 작성

2줄 딸깍이다. 

 

여기서 문제! 중간에 이런 코드가 있으면 돌아갈까? 안 돌아갈까? 

t = torch.randint(self.min_step, self.max_step, (1,))

정답은 안 돌아간다. 왜냐면 새로운 torch tensor를 만들었는데, to(device)를 안했기 떄문에 cpu에 올라간다. 

때문에 .cuda or .to(device)를 해줘야 한다. 

 

self.scheduler.alphas_cumprod:

  • 이 부분은 노이즈 스케줄러에서 사용하는 alphas_cumprod 값입니다. 확산 모델(예: DDPM)에서, alpha 값은 각 단계에서 노이즈가 어떻게 추가되거나 제거되는지를 나타내는 중요한 변수입니다.
  • alphas_cumprod는 누적된 alpha 값으로, 각 시간 단계에서 노이즈가 추가된 후의 변화를 나타냅니다. 이 값은 이미지가 노이즈를 제거하거나 추가하는 과정에서 사용됩니다.

만약 모델이 노이즈를 점진적으로 제거하면서 이미지를 복원하는 작업을 한다면, self.alphas 값은 각 스텝마다 그 과정에서의 노이즈를 조정하는 데 사용됩니다. alphas_cumprod는 각 스텝에서의 노이즈 비율이기 때문에 이 값을 기반으로 이미지 생성에 필요한 노이즈 패턴을 정교하게 제어할 수 있습니다.

 

음? 이게 확률 분포와 관련된 값인가? 

그 뒤에서 self.schedular를 통해서 노이즈를 더해주게 된다. 

 

아래 코드는 뭐하는 코드일까? latents와 동일한 벡터 크기의 노이즈를 만들어준다. 

그리고 스케쥴러에 잠재 벡터와 노이즈를 추가한다. 아마 저 스케쥴러에는 확률 분포를 만들어주는 뭔가가 있을 것이다. 

  • self.scheduler.add_noise(latents, noise, t): latents에 노이즈를 추가하는 과정입니다. 이 작업은 확산 과정에서 이미지에 노이즈를 더하는 것과 동일합니다. t는 시간 단계(time step)로, 노이즈가 얼마나 추가될지를 결정합니다
noise = torch.randn_like(latents)
latents_noisy = self.scheduler.add_noise(latents, noise, t)
  • self.unet(...): UNet은 노이즈를 예측하는 모델입니다. 여기서는 이미 학습된 UNet이 사용되고 있으며, latents_noisy를 입력받아 노이즈를 예측합니다. 

결국 노이즈를 만들고 u-net 구조를 통해서 예측하는 것이다. 

encoder_hidden_states로 들어가는 부분이 아마 condition이라고 생각된다. 그렇다면 u-net 구조는 원래 그것을 상정하고 만든건가? 

latent_model_input = torch.cat([latents_noisy] * 2)
tt = torch.cat([t] * 2)
noise_pred = self.unet(latent_model_input, tt, encoder_hidden_states=text_embeddings).sample

그 아래 코드에서는 noise_pred에 대해서 chunk라는 매서드를 사용한다. 

noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_pos - noise_pred_uncond)

몰라서 chatgpt한테 물어봤다.

chunk는 PyTorch 텐서에서 특정한 크기로 텐서를 분할할 때 사용하는 메서드입니다. 주어진 텐서를 일정한 크기로 나누어 여러 개의 작은 텐서로 나눕니다. noise_pred.chunk(2)는 noise_pred 텐서를 2개로 나누겠다는 의미입니다. 이렇게 chunk는 텐서를 여러 개로 나누어, 네트워크의 여러 부분에서 데이터를 분리하거나 병렬 연산을 쉽게 수행할 수 있게 해줍니다.

아! 아까 cat 해줬던 것을 다시 나누는 모양이다. 

그래도 모르겠어서 다시 물어봤다. 이 코드는 클래스 조건 생성(class-conditional generation)에서 사용하는 가이던스 스케일링(guidance scaling) 기법과 관련이 있습니다. 특히, 클래스 없는(Unconditional) 예측과 클래스 조건부(Conditional) 예측을 결합해, 원하는 특성을 강조하여 이미지를 생성하는 방식입니다. 아하 condition을 위한 코드인가 보다. 

  • 이 식은 가이던스 스케일링(guidance scaling)을 적용하여 최종적인 노이즈 예측 값을 계산하는 부분입니다.
  • 가이던스 스케일링은 조건부 예측(noise_pred_pos)과 조건 없는 예측(noise_pred_uncond) 사이의 차이를 조정하여, 원하는 조건(예: 특정 클래스, 이미지 속성 등)이 더 잘 반영된 결과를 얻기 위한 방법입니다.

그 아래에는 또 간단한 코드가 나오는데 이것도 이해하지 못해서 물어봤다. 

# w(t), sigma_t^2
w = (1 - self.alphas[t])

이 코드는 딥러닝 모델에서 주로 사용하는 노이즈 스케줄링(noise scheduling) 또는 디노이징(diffusion) 과정과 관련이 있습니다.

w = (1 - self.alphas[t])는 특정 시점 t에서 노이즈나 가중치를 조절하는 데 사용되는 변수 w를 계산하는 부분입니다.

노이즈 스케줄링과 alphas

  • 알파(alphas)는 시간 스텝에 따라 노이즈가 어떻게 변화하는지를 결정합니다.
  • 보통 노이즈 추가 과정에서는 시간이 지날수록 더 많은 노이즈가 추가되고, 디노이징 과정에서는 노이즈가 점차 제거됩니다. 이 때 알파 계수는 각 시간 스텝에서 노이즈와 신호 사이의 비율을 조절하는 역할을 합니다.
  • 알파가 1에 가까울수록 신호가 더 유지되고, 0에 가까울수록 노이즈가 더 많이 추가된 상태를 나타냅니다.

w = (1 - self.alphas[t])의 역할

  • w는 시간 t에서 노이즈의 분산 또는 가중치를 계산하는 데 사용됩니다.
    • 노이즈의 분산은 일반적으로 1 - 알파로 계산되며, 이는 노이즈가 시간이 지남에 따라 어떻게 변화하는지를 나타냅니다.
    • 노이즈 스케일링을 할 때, w는 각 시간 스텝에서 현재 이미지에 추가될 노이즈의 양을 제어합니다.

이제 정의에 따라서 loss를 계산하자. 

 

grad라는 명칭이 중요하진 않다고 생각한다. 

저기 있는 식을 먼저 구현하자.

grad = (noise_pred - noise)

그 다음은 아까 구했던 w값을 곱해주자. 

grad *= w

 

torch.nan_to_num(grad)는 PyTorch에서 사용되는 함수로, 주어진 텐서 grad 내에서 NaN (Not a Number), 무한대(Infinity)와 같은 특수한 값을 안전하게 다른 값으로 대체하는 기능을 제공합니다.

grad = torch.nan_to_num(grad)

음... 근데 지금 의문이 하나 드는게 결국 mse를 통해서 계산해야하는 것은 저 noise_pred와 noise의 차이다. 

그런데 과제에서는 아래와 같이 진행한다. 

아래 과제에서는 grad에 저장했다가 target에 넣었다가 다시 latent와 비교하는데 

결국 이게 Noise_pred랑 noise랑 비교하는 것이다. 

        grad = w * (noise_pred - noise)   # None 지우고 코드 작성

        #############################################################

        grad = torch.nan_to_num(grad)

        #############################################################
        # TODO: SDS loss에서 latents와의 비교를 위한 targets을 작성하세요.

        targets = (latents - grad)  # None 지우고 코드 작성

        #############################################################
        targets = targets.detach()
        loss = 0.5 * F.mse_loss(latents.float(), targets, reduction='sum') / latents.shape[0]

아무튼 결국 이것도 수식을 따르는 식이다. 

loss를 반환하면 마무리다. 

728x90
저작자표시