From 2117ef3295fb14ead27dd1b51ae79976d1f84ffc Mon Sep 17 00:00:00 2001
From: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com>
Date: Wed, 15 Jan 2025 17:10:10 +0100
Subject: [PATCH] Add pixart sigma and fix NOTthing (#762)

# What does this PR do?

as per title
---
 docs/source/inference_tutorials/stable_diffusion.mdx |  5 ++---
 docs/source/package_reference/supported_models.mdx   |  1 +
 optimum/neuron/__init__.py                           |  2 ++
 optimum/neuron/modeling_diffusion.py                 | 10 ++++++++++
 tests/cli/test_export_cli.py                         |  2 +-
 5 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/docs/source/inference_tutorials/stable_diffusion.mdx b/docs/source/inference_tutorials/stable_diffusion.mdx
index 01ff55394..f42e9340e 100644
--- a/docs/source/inference_tutorials/stable_diffusion.mdx
+++ b/docs/source/inference_tutorials/stable_diffusion.mdx
@@ -48,7 +48,7 @@ optimum-cli export neuron --model stabilityai/stable-diffusion-2-1-base \
   --batch_size 1 \
   --height 512 `# height in pixels of generated image, eg. 512, 768` \
   --width 512 `# width in pixels of generated image, eg. 512, 768` \
-  --num_images_per_prompt 4 `# number of images to generate per prompt, defaults to 1` \
+  --num_images_per_prompt 1 `# number of images to generate per prompt, defaults to 1` \
   --auto_cast matmul `# cast only matrix multiplication operations` \
   --auto_cast_type bf16 `# cast operations from FP32 to BF16` \
   sd_neuron/
@@ -231,7 +231,7 @@ optimum-cli export neuron --model stabilityai/stable-diffusion-xl-base-1.0 \
   --batch_size 1 \
   --height 1024 `# height in pixels of generated image, eg. 768, 1024` \
   --width 1024 `# width in pixels of generated image, eg. 768, 1024` \
-  --num_images_per_prompt 4 `# number of images to generate per prompt, defaults to 1` \
+  --num_images_per_prompt 1 `# number of images to generate per prompt, defaults to 1` \
   --auto_cast matmul `# cast only matrix multiplication operations` \
   --auto_cast_type bf16 `# cast operations from FP32 to BF16` \
   sd_neuron_xl/
@@ -340,7 +340,6 @@ prompt = "A majestic lion jumping from a big stone at night"
 base = NeuronStableDiffusionXLPipeline.from_pretrained("sd_neuron_xl/")
 image = base(
     prompt=prompt,
-    num_images_per_prompt=num_images_per_prompt,
     num_inference_steps=40,
     denoising_end=0.8,
     output_type="latent",
diff --git a/docs/source/package_reference/supported_models.mdx b/docs/source/package_reference/supported_models.mdx
index d44c2e613..f0dde3278 100644
--- a/docs/source/package_reference/supported_models.mdx
+++ b/docs/source/package_reference/supported_models.mdx
@@ -75,6 +75,7 @@ limitations under the License.
 | SDXL Turbo                    | text-to-image, image-to-image, inpaint                                                                                                 |
 | LCM                           | text-to-image                                                                                                                          |
 | PixArt-α                      | text-to-image                                                                                                                          |
+| PixArt-Σ                      | text-to-image                                                                                                                          |
 
 ## Sentence Transformers
 
diff --git a/optimum/neuron/__init__.py b/optimum/neuron/__init__.py
index 6900143d9..471befb7a 100644
--- a/optimum/neuron/__init__.py
+++ b/optimum/neuron/__init__.py
@@ -60,6 +60,7 @@
         "NeuronStableDiffusionControlNetPipeline",
         "NeuronStableDiffusionXLControlNetPipeline",
         "NeuronPixArtAlphaPipeline",
+        "NeuronPixArtSigmaPipeline",
     ],
     "modeling_decoder": ["NeuronDecoderModel"],
     "modeling_seq2seq": ["NeuronModelForSeq2SeqLM"],
@@ -98,6 +99,7 @@
         NeuronDiffusionPipelineBase,
         NeuronLatentConsistencyModelPipeline,
         NeuronPixArtAlphaPipeline,
+        NeuronPixArtSigmaPipeline,
         NeuronStableDiffusionControlNetPipeline,
         NeuronStableDiffusionImg2ImgPipeline,
         NeuronStableDiffusionInpaintPipeline,
diff --git a/optimum/neuron/modeling_diffusion.py b/optimum/neuron/modeling_diffusion.py
index b3bec452c..53f2df76e 100644
--- a/optimum/neuron/modeling_diffusion.py
+++ b/optimum/neuron/modeling_diffusion.py
@@ -78,6 +78,7 @@
         LatentConsistencyModelPipeline,
         LCMScheduler,
         PixArtAlphaPipeline,
+        PixArtSigmaPipeline,
         StableDiffusionControlNetPipeline,
         StableDiffusionImg2ImgPipeline,
         StableDiffusionInpaintPipeline,
@@ -1504,6 +1505,15 @@ def __init__(self, **kwargs):
         self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor)
 
 
+class NeuronPixArtSigmaPipeline(NeuronDiffusionPipelineBase, PixArtSigmaPipeline):
+    main_input_name = "prompt"
+    auto_model_class = PixArtSigmaPipeline
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+
 class NeuronStableDiffusionXLPipeline(
     NeuronStableDiffusionXLPipelineMixin, NeuronDiffusionPipelineBase, StableDiffusionXLPipeline
 ):
diff --git a/tests/cli/test_export_cli.py b/tests/cli/test_export_cli.py
index 6ab40fcea..89dde7d9e 100644
--- a/tests/cli/test_export_cli.py
+++ b/tests/cli/test_export_cli.py
@@ -159,7 +159,7 @@ def test_stable_diffusion(self):
                         "--width",
                         "64",
                         "--num_images_per_prompt",
-                        "4",
+                        "1",
                         "--auto_cast",
                         "matmul",
                         "--auto_cast_type",