test

huggingface · Jan 15, 2025 · 10903a7 · 10903a7
1 parent cc3d5d6
commit 10903a7
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 6 deletions.
diff --git 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
@@ -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/exporters/neuron/convert.py b/optimum/exporters/neuron/convert.py
@@ -557,8 +557,7 @@ def export_neuronx(
     compiler_args.extend(["--optlevel", optlevel])
     logger.info(f"Using Neuron: --optlevel {optlevel}")
 
-    if getattr(config._config, "is_encoder_decoder", False):
-        compiler_args.extend(["--model-type", "transformer"])
+    compiler_args.extend(["--model-type", "transformer"])
 
     compiler_args = add_stable_diffusion_compiler_args(config, compiler_args)  # diffusers specific
 

diff --git 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
@@ -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
@@ -159,7 +159,7 @@ def test_stable_diffusion(self):
                         "--width",
                         "64",
                         "--num_images_per_prompt",
-                        "4",
+                        "1",
                         "--auto_cast",
                         "matmul",
                         "--auto_cast_type",