디퓨전 모델

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch

# Stable Diffusion XL 파이프라인 로드
pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
)
pipe = pipe.to("cuda")

# 빠른 스케줄러로 교체 (20스텝으로 고품질 생성)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config
)

# 텍스트-to-이미지 생성
prompt = "A futuristic city skyline at sunset, cyberpunk style, 8k, detailed"
negative_prompt = "blurry, low quality, distorted"

image = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    num_inference_steps=25,
    guidance_scale=7.5,  # Classifier-Free Guidance 강도
    height=1024,
    width=1024
).images[0]

image.save("generated_city.png")

# 이미지-to-이미지 변환 (기존 이미지 스타일 변경)
from diffusers import AutoPipelineForImage2Image
from PIL import Image

img2img_pipe = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/sdxl-turbo",
    torch_dtype=torch.float16
).to("cuda")

init_image = Image.open("input.jpg").resize((512, 512))
result = img2img_pipe(
    prompt="oil painting style",
    image=init_image,
    strength=0.7,  # 원본 유지 정도 (0=원본, 1=완전 새로 생성)
    num_inference_steps=4  # SDXL Turbo는 4스텝으로 충분
).images[0]

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleDiffusion:
    """DDPM의 핵심 개념을 단순화한 구현"""
    def __init__(self, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.timesteps = timesteps
        # 노이즈 스케줄 정의
        self.betas = torch.linspace(beta_start, beta_end, timesteps)
        self.alphas = 1.0 - self.betas
        self.alpha_cumprod = torch.cumprod(self.alphas, dim=0)

    def forward_diffusion(self, x0, t):
        """Forward process: 원본에 노이즈 추가"""
        noise = torch.randn_like(x0)
        sqrt_alpha_cumprod = self.alpha_cumprod[t].sqrt()
        sqrt_one_minus = (1 - self.alpha_cumprod[t]).sqrt()

        # x_t = sqrt(alpha_cumprod) * x_0 + sqrt(1-alpha_cumprod) * noise
        xt = sqrt_alpha_cumprod.view(-1, 1, 1, 1) * x0 + \
             sqrt_one_minus.view(-1, 1, 1, 1) * noise
        return xt, noise

    def training_step(self, model, x0, optimizer):
        """학습: 노이즈 예측"""
        batch_size = x0.shape[0]
        t = torch.randint(0, self.timesteps, (batch_size,))

        xt, noise = self.forward_diffusion(x0, t)
        predicted_noise = model(xt, t)

        loss = F.mse_loss(predicted_noise, noise)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        return loss.item()

    @torch.no_grad()
    def sample(self, model, shape, device):
        """Reverse process: 노이즈에서 이미지 생성"""
        x = torch.randn(shape, device=device)

        for t in reversed(range(self.timesteps)):
            t_batch = torch.full((shape[0],), t, device=device)
            predicted_noise = model(x, t_batch)

            alpha = self.alphas[t]
            alpha_cumprod = self.alpha_cumprod[t]
            beta = self.betas[t]

            if t > 0:
                noise = torch.randn_like(x)
            else:
                noise = 0

            # 노이즈 제거 스텝
            x = (1 / alpha.sqrt()) * (
                x - (beta / (1 - alpha_cumprod).sqrt()) * predicted_noise
            ) + beta.sqrt() * noise

        return x