Update infrance_text_pop.py

infrance_text_pop.py

@@ -128,9 +128,11 @@ if __name__ == '__main__':
     import torch.nn as nn
 
     # Initialize the model
-    pl_checkpoint = '/Users/mac/Downloads/model-step-step=35000.ckpt'  # Replace with your actual checkpoint path
+    checkpoint = 'VARtext_v1.pth'  # Replace with your actual checkpoint path
     device = 'mps'
-    model = InrenceTextVAR(pl_checkpoint=pl_checkpoint, device=device)
+    model = InrenceTextVAR(device=device)
+    state_dict = torch.load(checkpoint,map_location = "cpu")
+    model.load_state_dict(state_dict)
     model.to(device)
 
     def generate_image_gradio(text, beta=1.0, seed=None, more_smooth=False, top_k=0, top_p=0.9):