Open-Sora 모델 아키텍처 심층 분석: STDiT3, VAE, T5의 완벽 이해

@MODELS.register_module("text_embedder")
class HFEmbedder(nn.Module):
    def __init__(self, from_pretrained: str, max_length: int, shardformer: bool = False, **hf_kwargs):
        super().__init__()
        self.is_clip = "openai" in from_pretrained
        self.max_length = max_length
        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"

        if self.is_clip:
            self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(from_pretrained, max_length=max_length)
            self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(from_pretrained, **hf_kwargs)
        else:
            self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(
                from_pretrained, max_length=max_length, legacy=True
            )
            self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(from_pretrained, **hf_kwargs)

def shardformer_t5(t5: T5EncoderModel) -> T5EncoderModel:
    """T5 모델을 샤딩하여 최적화"""
    dtype = t5.shared.weight.dtype
    shard_config = ShardConfig(
        enable_tensor_parallelism=False,  # 텐서 병렬화 비활성화
        enable_jit_fused=True,           # JIT 융합 최적화 활성화
    )
    shard_former = ShardFormer(shard_config=shard_config)
    optim_model, _ = shard_former.optimize(t5, policy=T5EncoderPolicy())
    return optim_model.to(dtype).eval().requires_grad_(False)

def forward(self, text: list[str], added_tokens: int = 0, seq_align: int = 1) -> Tensor:
    # 1. 토크나이징
    batch_encoding = self.tokenizer(
        text,
        truncation=True,
        max_length=self.max_length,
        return_length=False,
        return_overflowing_tokens=False,
        padding="max_length",
        return_tensors="pt",
    )
    
    # 2. 시퀀스 정렬 (배치 처리 최적화)
    seq_len = batch_encoding["input_ids"].shape[1]
    if (added_tokens + seq_len) % seq_align != 0:
        num_pad_tokens = seq_align - (added_tokens + seq_len) % seq_align
        batch_encoding["input_ids"] = nn.functional.pad(
            batch_encoding["input_ids"], (0, num_pad_tokens), value=self.tokenizer.pad_token_id
        )

    # 3. 임베딩 생성
    outputs = self.hf_module(
        input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
        attention_mask=None,
        output_hidden_states=False,
    )
    return outputs[self.output_key]

@dataclass
class AutoEncoderConfig:
    from_pretrained: str | None
    cache_dir: str | None
    resolution: int          # 해상도
    in_channels: int         # 입력 채널 수
    ch: int                  # 기본 채널 수
    out_ch: int             # 출력 채널 수
    ch_mult: list[int]      # 채널 배수 (다운샘플링용)
    num_res_blocks: int     # 잔차 블록 수
    z_channels: int         # 잠재 공간 채널 수
    scale_factor: float     # 스케일 팩터
    shift_factor: float     # 시프트 팩터
    sample: bool = True     # 샘플링 여부

class AttnBlock(nn.Module):
    def __init__(self, in_channels: int):
        super().__init__()
        self.in_channels = in_channels
        
        # Group Normalization
        self.norm = nn.GroupNorm(
            num_groups=32, 
            num_channels=in_channels, 
            eps=1e-6, 
            affine=True
        )
        
        # Query, Key, Value 프로젝션
        self.q = nn.Conv2d(in_channels, in_channels, 1, 1, 0)
        self.k = nn.Conv2d(in_channels, in_channels, 1, 1, 0)
        self.v = nn.Conv2d(in_channels, in_channels, 1, 1, 0)
        
        # 출력 프로젝션
        self.proj_out = nn.Conv2d(in_channels, in_channels, 1, 1, 0)

class DiagonalGaussianDistribution:
    def __init__(self, parameters: Tensor, deterministic: bool = False):
        self.parameters = parameters
        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
        self.deterministic = deterministic
        self.std = torch.exp(0.5 * self.logvar)
        self.var = torch.exp(self.logvar)

    def sample(self) -> Tensor:
        if self.deterministic:
            return self.mean
        else:
            return self.mean + self.std * torch.randn_like(self.mean)

# Gradio 앱에서의 STDiT3 로딩
from opensora.models.stdit.stdit3 import STDiT3

model_kwargs = {k: v for k, v in config.model.items() 
                if k not in ("type", "from_pretrained", "force_huggingface")}

print("Load STDIT3 from ", weight_path)
stdit = STDiT3.from_pretrained(weight_path, **model_kwargs).cuda()

# 해상도별 모델 가중치
HF_STDIT_MAP = {
    "t2v": {
        "360p": "hpcaitech/OpenSora-STDiT-v4-360p",
        "720p": "hpcaitech/OpenSora-STDiT-v4",
    },
    "i2v": "hpcaitech/OpenSora-STDiT-v4-i2v",
}

# 샘플링 과정
samples = scheduler.sample(
    stdit,                    # STDiT3 모델
    text_encoder,            # T5 텍스트 인코더
    z=z,                     # 초기 노이즈
    z_cond=ref,              # 조건부 잠재 벡터
    z_cond_mask=x_cond_mask, # 조건부 마스크
    prompts=batch_prompts_loop,  # 배치 프롬프트
    device=device,
    additional_args=model_args,
    progress=True,
    mask=masks,
    mask_index=mask_index,
    image_cfg_scale=image_cfg_scale,
    use_sdedit=use_sdedit,
    use_oscillation_guidance_for_text=use_oscillation_guidance_for_text,
    use_oscillation_guidance_for_image=use_oscillation_guidance_for_image,
)

def build_models(mode, resolution, enable_optimization=False):
    """모델 구축 및 초기화"""
    
    # 1. VAE 구축
    vae = build_module(config.vae, MODELS).cuda()
    
    # 2. 텍스트 인코더 구축
    text_encoder = build_module(config.text_encoder, MODELS)
    text_encoder.t5.model = text_encoder.t5.model.cuda()
    
    # 3. STDiT3 구축
    stdit = STDiT3.from_pretrained(weight_path, **model_kwargs).cuda()
    
    # 4. 스케줄러 구축
    scheduler = build_module(config.scheduler, SCHEDULERS)
    
    # 5. CFG를 위한 임베더 연결
    text_encoder.y_embedder = stdit.y_embedder
    
    # 6. 최적화 및 평가 모드
    vae = vae.to(torch.bfloat16).eval()
    text_encoder.t5.model = text_encoder.t5.model.eval()
    stdit = stdit.to(torch.bfloat16).eval()
    
    return vae, text_encoder, stdit, scheduler, config

# 참조 이미지 처리
refs = collect_references_batch(refs, vae, image_size)

# 조건부 참조 준비
if mode == "i2v":
    image_cfg_scale = config.get("image_cfg_scale", 7.5)
    ref, mask_index = prep_ref_and_mask(
        cond_type, condition_frame_length, refs, target_shape, 
        num_loop, device, dtype
    )

# 모델을 bfloat16으로 변환
vae = vae.to(torch.bfloat16).eval()
stdit = stdit.to(torch.bfloat16).eval()

# T5는 fp32 유지 (정확도 보장)
text_encoder.t5.model = text_encoder.t5.model.eval()

# 멀티 GPU 추론
torchrun --nproc_per_node 8 --standalone scripts/diffusion/inference.py \
    configs/diffusion/inference/t2i2v_768px.py \
    --save-dir samples --prompt "raining, sea"

# CUDA 메모리 정리
torch.cuda.empty_cache()

# 그래디언트 비활성화
with torch.inference_mode():
    # 추론 코드
    pass

# 해상도별 점진적 학습
resolutions = ["144p", "360p", "720p"]
for resolution in resolutions:
    train_model(resolution)

# 소셜 미디어 숏폼 생성
prompt = "A cat playing with a ball in a sunny garden, cute and playful"
video = generate_video(prompt, aspect_ratio="9:16", length="5s")

# 기존 이미지에서 비디오 생성
reference_image = load_image("portrait.jpg")
prompt = "The person smiles and nods gently"
video = generate_video(prompt, reference_image=reference_image)

# 광고 컨셉 시각화
prompt = "Modern minimalist product showcase, elegant lighting"
video = generate_video(prompt, resolution="720p", motion_score=3)