reestore

huggingface · Jul 25, 2024 · f7aad56 · f7aad56
1 parent ecdc197
commit f7aad56
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 7 deletions.
diff --git a/optimum/exporters/neuron/convert.py b/optimum/exporters/neuron/convert.py
@@ -352,13 +352,6 @@ def export_models(
         output_path.parent.mkdir(parents=True, exist_ok=True)
 
         try:
-            # TODO: Remove after the weights/neff separation compilation of sdxl is patched by a neuron sdk release: https://github.com/aws-neuron/aws-neuron-sdk/issues/859
-            if not inline_weights_to_neff and getattr(sub_neuron_config, "is_sdxl", False):
-                logger.warning(
-                    "The compilation of SDXL's unet with the weights/neff separation is broken since the Neuron sdk 2.18 release. `inline_weights_to_neff` will be set to True and the caching will be disabled. If you still want to separate the neff and weights, please downgrade your Neuron setup to the 2.17.1 release."
-                )
-                inline_weights_to_neff = True
-
             start_time = time.time()
             neuron_inputs, neuron_outputs = export(
                 model=submodel,

diff --git a/tests/cli/test_export_cli.py b/tests/cli/test_export_cli.py
@@ -250,6 +250,38 @@ def test_stable_diffusion_single_controlnet(self):
 
     @requires_neuronx
     def test_stable_diffusion_xl(self):
+        model_id = "echarlaix/tiny-random-stable-diffusion-xl"
+        with tempfile.TemporaryDirectory() as tempdir:
+            subprocess.run(
+                [
+                    "optimum-cli",
+                    "export",
+                    "neuron",
+                    "--model",
+                    model_id,
+                    "--inline-weights-neff",
+                    "--task",
+                    "stable-diffusion-xl",
+                    "--batch_size",
+                    "1",
+                    "--height",
+                    "64",
+                    "--width",
+                    "64",
+                    "--num_images_per_prompt",
+                    "4",
+                    "--auto_cast",
+                    "matmul",
+                    "--auto_cast_type",
+                    "bf16",
+                    tempdir,
+                ],
+                shell=False,
+                check=True,
+            )
+
+    @requires_neuronx
+    def test_stable_diffusion_xl_non_inlined(self):
         model_id = "echarlaix/tiny-random-stable-diffusion-xl"
         with tempfile.TemporaryDirectory() as tempdir:
             subprocess.run(