Update pipeline.py

pipeline.py

@@ -7,7 +7,7 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 import os
 from .vae import AutoencoderKL
-from .mar import mar
+from .mar import mar_base, mar_large, mar_huge
 
 # inheriting from DiffusionPipeline for HF
 class MARModel(DiffusionPipeline):  
@@ -33,12 +33,27 @@ class MARModel(DiffusionPipeline):
         model_type = kwargs.get("model_type", "mar_base")
 
 
-        self.model = mar.__dict__[model_type](
-            buffer_size=buffer_size,
-            diffloss_d=diffloss_d,
-            diffloss_w=diffloss_w,
-            num_sampling_steps=str(num_sampling_steps)
-        ).to(device)
+        if model_type == "mar_base":
+            self.model = mar_base(
+                buffer_size=buffer_size,
+                diffloss_d=diffloss_d,
+                diffloss_w=diffloss_w,
+                num_sampling_steps=str(num_sampling_steps)
+            ).to(device)
+        elif model_type == "mar_large":
+            self.model = mar_large(
+                buffer_size=buffer_size,
+                diffloss_d=diffloss_d,
+                diffloss_w=diffloss_w,
+                num_sampling_steps=str(num_sampling_steps)
+            ).to(device)
+        elif model_type == "mar_huge":
+            self.model = mar_huge(
+                buffer_size=buffer_size,
+                diffloss_d=diffloss_d,
+                diffloss_w=diffloss_w,
+                num_sampling_steps=str(num_sampling_steps)
+            ).to(device)
         # download and load the model weights (.safetensors or .pth)
         model_checkpoint_path = hf_hub_download(
             repo_id=kwargs.get("repo_id", "jadechoghari/mar"),