add safetensors

pipeline.py

@@ -6,8 +6,7 @@ import sys
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 import os
-from torchvision.utils import save_image
-from PIL import Image
+from safetensors.torch import load_file
 from .vae import AutoencoderKL
 from .mar import mar_base, mar_large, mar_huge
 
@@ -46,20 +45,22 @@ class MARModel(DiffusionPipeline):
         if model_type == "mar_base":
             diffloss_d = 6
             diffloss_w = 1024
+            model_path = "mar-base.safetensors"
         elif model_type == "mar_large":
             diffloss_d = 8
             diffloss_w = 1280
+            model_path = "mar-large.safetensors"
         elif model_type == "mar_huge":
             diffloss_d = 12
             diffloss_w = 1536
+            model_path = "mar-huge.safetensors"
         else:
             raise NotImplementedError
         # download and load the model weights (.safetensors or .pth)
         model_checkpoint_path = hf_hub_download(
             repo_id=kwargs.get("repo_id", "jadechoghari/mar"),
-            filename=kwargs.get("model_filename", "checkpoint-last.pth")
+            filename=kwargs.get("model_filename", model_path)
         )
-        model_checkpoint_path = kwargs.get("model_checkpoint_path", "./mar/checkpoint-last.pth")
 
         model_fn = model_mapping[model_type]
 
@@ -70,7 +71,8 @@ class MARModel(DiffusionPipeline):
             num_sampling_steps=str(num_sampling_steps_diffloss)
         ).cuda()
 
-        state_dict = torch.load(f"./mar/checkpoint-last.pth")["model_ema"]
+        # use safetensors
+        state_dict = load_file(safetensors_path)
         model.load_state_dict(state_dict)
         model.eval()
 
@@ -85,7 +87,7 @@ class MARModel(DiffusionPipeline):
         vae = vae.to(device).eval()
 
         # set up user-specified or default values for generation
-        seed = kwargs.get("seed", 0)
+        seed = kwargs.get("seed", 6)
         torch.manual_seed(seed)
         np.random.seed(seed)
 
@@ -93,9 +95,7 @@ class MARModel(DiffusionPipeline):
         cfg_scale = kwargs.get("cfg_scale", 4)
         cfg_schedule = kwargs.get("cfg_schedule", "constant")
         temperature = kwargs.get("temperature", 1.0)
-        # class_labels = kwargs.get("class_labels", 207, 360, 388, 113, 355, 980, 323, 979)
-        class_labels = 207, 360, 388, 113, 355, 980, 323, 979
-        print("the labels", class_labels)
+        class_labels = kwargs.get("class_labels", 207, 360, 388, 113, 355, 980, 323, 979)
 
         # generate the tokens and images
         with torch.cuda.amp.autocast():
@@ -113,7 +113,7 @@ class MARModel(DiffusionPipeline):
     
         # save the images
         image_path = os.path.join(output_dir, "sampled_image.png")
-        samples_per_row = kwargs.get("samples_per_row", 6)
+        samples_per_row = kwargs.get("samples_per_row", 4)
     
         save_image(
             sampled_images, image_path, nrow=int(samples_per_row), normalize=True, value_range=(-1, 1)