| 22 | |
| 23 | @dataclass |
| 24 | class T2IConfig: |
| 25 | |
| 26 | model_path: str = "../checkpoint" |
| 27 | device: str = "cuda:1" |
| 28 | dtype: torch.dtype = torch.bfloat16 |
| 29 | scale: int = 0 |
| 30 | seq_len: int = 729 |
| 31 | top_p: float = 0.95 |
| 32 | top_k: int = 1200 |
| 33 | |
| 34 | class TextToImageInference: |
| 35 | def __init__(self, config: T2IConfig): |