File size: 402 Bytes
14ce5a9
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
from .base_model import BaseModel
from .ema_model import EMAModel
from .losses import (
    ReconstructionLoss_Stage1,
    ReconstructionLoss_Stage2,
    ReconstructionLoss_Single_Stage,
    MLMLoss,
    ARLoss,
)
from .blocks import TiTokEncoder, TiTokDecoder, TATiTokDecoder, UViTBlock
from .maskgit_vqgan import Decoder as Pixel_Decoder
from .maskgit_vqgan import VectorQuantizer as Pixel_Quantizer