Update handler.py

handler.py

@@ -13,45 +13,52 @@ def get_default_args():
     parser = argparse.ArgumentParser()
     
     # Model configuration
-    parser.add_argument("--model", type=str, default="HYVideo-T/2")
+    parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill")
     parser.add_argument("--model-resolution", type=str, default="720p", choices=["540p", "720p"])
-    parser.add_argument("--latent-channels", type=int, default=16)  # Changed from 4 to match VAE
+    parser.add_argument("--latent-channels", type=int, default=16)
     parser.add_argument("--precision", type=str, default="bf16", choices=["bf16", "fp32", "fp16"])
-    parser.add_argument("--rope-theta", type=int, default=256)  # Changed to match original config
-    parser.add_argument("--load-key", type=str, default="module", help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.")
+    parser.add_argument("--rope-theta", type=int, default=256)
+    parser.add_argument("--load-key", type=str, default="module")
     
     # VAE settings
     parser.add_argument("--vae", type=str, default="884-16c-hy")
-    parser.add_argument("--vae-precision", type=str, default="fp16")  # Changed to fp16 to match docs
-    parser.add_argument("--vae-tiling", action="store_true", default=True)  # Set default to True
+    parser.add_argument("--vae-precision", type=str, default="fp16")
+    parser.add_argument("--vae-tiling", action="store_true", default=True)
     
     # Text encoder settings
-    parser.add_argument("--text-encoder", type=str, default="llm")  # Changed to llm
-    parser.add_argument("--text-encoder-precision", type=str, default="fp16")  # Changed to fp16
-    parser.add_argument("--text-states-dim", type=int, default=4096)  # Updated to match docs
-    parser.add_argument("--text-len", type=int, default=256)  # Updated to match docs
-    parser.add_argument("--tokenizer", type=str, default="llm")  # Changed to llm
+    parser.add_argument("--text-encoder", type=str, default="llm")
+    parser.add_argument("--text-encoder-precision", type=str, default="fp16")
+    parser.add_argument("--text-states-dim", type=int, default=4096)
+    parser.add_argument("--text-len", type=int, default=256)
+    parser.add_argument("--tokenizer", type=str, default="llm")
     
     # Prompt template settings
     parser.add_argument("--prompt-template", type=str, default="dit-llm-encode")
     parser.add_argument("--prompt-template-video", type=str, default="dit-llm-encode-video")
     
     # Additional text encoder settings
-    parser.add_argument("--hidden-state-skip-layer", type=int, default=2)  # Updated to match docs
+    parser.add_argument("--hidden-state-skip-layer", type=int, default=2)
     parser.add_argument("--apply-final-norm", action="store_true")
     parser.add_argument("--text-encoder-2", type=str, default="clipL")
-    parser.add_argument("--text-encoder-precision-2", type=str, default="fp16")  # Changed to fp16
-    parser.add_argument("--text-states-dim-2", type=int, default=768)  # Updated to match docs
+    parser.add_argument("--text-encoder-precision-2", type=str, default="fp16")
+    parser.add_argument("--text-states-dim-2", type=int, default=768)
     parser.add_argument("--tokenizer-2", type=str, default="clipL")
-    parser.add_argument("--text-len-2", type=int, default=77)  # Updated to match docs
+    parser.add_argument("--text-len-2", type=int, default=77)
+    
+    # Model architecture settings
+    parser.add_argument("--hidden-size", type=int, default=1024)
+    parser.add_argument("--heads-num", type=int, default=16)
+    parser.add_argument("--layers-num", type=int, default=24)
+    parser.add_argument("--mlp-ratio", type=float, default=4.0)
+    parser.add_argument("--use-guidance-net", action="store_true", default=True)
     
     # Inference settings
-    parser.add_argument("--denoise-type", type=str, default="flow")  # Changed to flow
+    parser.add_argument("--denoise-type", type=str, default="flow")
     parser.add_argument("--flow-shift", type=float, default=7.0)
-    parser.add_argument("--flow-reverse", action="store_true", default=False)
+    parser.add_argument("--flow-reverse", action="store_true", default=True)
     parser.add_argument("--flow-solver", type=str, default="euler")
     parser.add_argument("--use-linear-quadratic-schedule", action="store_true")
-    parser.add_argument("--linear-schedule-end", type=int, default=25)  # Updated to match docs
+    parser.add_argument("--linear-schedule-end", type=int, default=25)
     
     # Hardware settings
     parser.add_argument("--use-cpu-offload", action="store_true", default=False)