Refactor argument parsing and improve error handling in model loading

Update .gitignore to exclude PNG files and remove unnecessary output images

.gitignore

@@ -1 +1,2 @@
 *.pyc
+*.png

inference.py

@@ -51,7 +51,6 @@ def main(args):
 
     # Load model configuration and model
     config = MMDiTConfig.from_json_file(args.model_config)
-    config.vae_type = args.vae_type  # VAE overriding
     config.height = args.resolution
     config.width = args.resolution
 
@@ -135,11 +134,10 @@ if __name__ == "__main__":
     # parser.add_argument("--slg", type=int, nargs="*", default=None, help="")
     parser.add_argument("--steps", type=int, default=50, help="Number of steps for image generation")
     parser.add_argument("--resolution", type=int, default=256, help="Resolution of output images")
-    parser.add_argument("--batch-size", type=int, default=32)
-    parser.add_argument("--streaming", action="store_true")
+    parser.add_argument("--batch-size", type=int, default=32,help="Batch size for image generation")
+    parser.add_argument("--streaming", action="store_true", help="Enable streaming mode for intermediate steps")
     parser.add_argument("--noisy-pad", action="store_true")
     parser.add_argument("--zero-masking", action="store_true")
-    parser.add_argument("--vae-type", type=str, default="SD3", help="Type of VAE")
     parser.add_argument("--prompt-file", type=str, default="prompt_128.txt", help="Path to the prompt file")
     parser.add_argument("--guidance-scales", type=float, nargs="*", default=None, help="List of guidance scales")
     parser.add_argument("--output-dir", type=str, default="output", help="Base output directory for generated images")

models/modeling_dit.py

@@ -13,7 +13,7 @@ try:
     MotifRMSNorm = motif_ops.T5LayerNorm
     ScaledDotProductAttention = None
     MotifFlashAttention = motif_ops.flash_attention
-except ImportError: # if motif_ops is not available
+except Exception: # if motif_ops is not available
     MotifRMSNorm = None
     ScaledDotProductAttention = None
     MotifFlashAttention = None