256

app.py

@@ -27,7 +27,7 @@ from diffusion.flow_matching import ODEEulerFlowMatchingSolver
 import utils
 import libs.autoencoder
 from libs.clip import FrozenCLIPEmbedder
-from configs import t2i_512px_clip_dimr
+from configs import t2i_512px_clip_dimr, t2i_256px_clip_dimr
 
 
 def unpreprocess(x: torch.Tensor) -> torch.Tensor:
@@ -93,7 +93,8 @@ def get_caption(llm: str, text_model, prompt_dict: dict, batch_size: int):
     return context, token_mask, tokens, captions
 
 # Load configuration and initialize models.
-config_dict = t2i_512px_clip_dimr.get_config()
+# config_dict = t2i_512px_clip_dimr.get_config()
+config_dict = t2i_256px_clip_dimr.get_config()
 config = ml_collections.ConfigDict(config_dict)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -108,7 +109,8 @@ MAX_IMAGE_SIZE = 1024  # Currently not used.
 
 # Load the main diffusion model.
 repo_id = "QHL067/CrossFlow"
-filename = "pretrained_models/t2i_512px_clip_dimr.pth"
+# filename = "pretrained_models/t2i_512px_clip_dimr.pth"
+filename = "pretrained_models/t2i_256px_clip_dimr.pth"
 checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
 nnet = utils.get_nnet(**config.nnet)
 nnet = nnet.to(device)