add cli test

huggingface · Dec 11, 2024 · 2e89b60 · 2e89b60
1 parent a0e4184
commit 2e89b60
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 5 deletions.
diff --git a/optimum/commands/export/neuronx.py b/optimum/commands/export/neuronx.py
@@ -116,9 +116,9 @@ def parse_args_neuronx(parser: "ArgumentParser"):
     optional_group.add_argument(
         "--torch_dtype",
         type=str,
-        default="auto",
-        choices=["auto", "bfloat16", "float16", "float32"],
-        help="Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the dtype will be automatically derived from the model's weights.",
+        default=None,
+        choices=["bfloat16", "float16", "float32"],
+        help="Override the default `torch.dtype` and load the model under this dtype. If `None` is passed, the dtype will be automatically derived from the model's weights.",
     )
     optional_group.add_argument(
         "--tensor_parallel_size",

diff --git a/optimum/exporters/neuron/model_configs/traced_configs.py b/optimum/exporters/neuron/model_configs/traced_configs.py
@@ -813,7 +813,6 @@ def patch_model_for_export(
         return super().patch_model_for_export(model=model, dummy_inputs=dummy_inputs, forward_with_tuple=True)
 
 
-@register_in_tasks_manager("t5-encoder", "text2text-generation")
 class T5EncoderBaseNeuronConfig(TextSeq2SeqNeuronConfig):
     ATOL_FOR_VALIDATION = 1e-3
     NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args(

diff --git a/tests/cli/test_export_cli.py b/tests/cli/test_export_cli.py
@@ -170,6 +170,36 @@ def test_stable_diffusion(self):
                     check=True,
                 )
 
+    @requires_neuronx
+    def test_pixart(self):
+        model_ids = ["hf-internal-testing/tiny-pixart-alpha-pipe"]
+        for model_id in model_ids:
+            with tempfile.TemporaryDirectory() as tempdir:
+                subprocess.run(
+                    [
+                        "optimum-cli",
+                        "export",
+                        "neuron",
+                        "--model",
+                        model_id,
+                        "--batch_size",
+                        "1",
+                        "--height",
+                        "8",
+                        "--width",
+                        "8",
+                        "--sequence_length",
+                        "16",
+                        "--num_images_per_prompt",
+                        "1",
+                        "--torch_dtype",
+                        "bfloat16",
+                        tempdir,
+                    ],
+                    shell=False,
+                    check=True,
+                )
+
     @requires_neuronx
     def test_stable_diffusion_multi_lora(self):
         model_id = "hf-internal-testing/tiny-stable-diffusion-torch"
@@ -196,7 +226,7 @@ def test_stable_diffusion_multi_lora(self):
                     lora_model_id,
                     "--lora_weight_names",
                     lora_weight_name,
-                    "lora_adapter_names",
+                    "--lora_adapter_names",
                     adpater_name,
                     "--lora_scales",
                     "0.9",