From f7aad565141cb926003b9cd57aab13fe90898ae6 Mon Sep 17 00:00:00 2001
From: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com>
Date: Thu, 25 Jul 2024 07:26:16 +0000
Subject: [PATCH] reestore

---
 optimum/exporters/neuron/convert.py |  7 -------
 tests/cli/test_export_cli.py        | 32 +++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/optimum/exporters/neuron/convert.py b/optimum/exporters/neuron/convert.py
index 810bbbedf..1b42db99b 100644
--- 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
index 72f84a50c..330f8e3f0 100644
--- 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(