[Stable Diffusion] Fix the image size value inferring (#167)

* fix the image size value inferring

* improve logging

* override push_to_hub as in decoder modeling

---------

Co-authored-by: Jingya Huang <[email protected]>

optimum/exporters/neuron/__main__.py

@@ -244,8 +244,7 @@ def main_export(
             )
 
             logger.info(
-                f"The {NEURON_COMPILER} export succeeded and the exported model was saved at: "
-                f"{output.parent.as_posix()}"
+                f"The {NEURON_COMPILER} export succeeded and the exported model was saved at: " f"{output.as_posix()}"
             )
         except ShapeError as e:
             raise e

optimum/exporters/neuron/convert.py

@@ -15,6 +15,7 @@
 """Neuron compiled model check and export functions."""
 
 import copy
+import time
 from collections import OrderedDict
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
@@ -289,6 +290,7 @@ def export_models(
 
     failed_models = []
     for i, model_name in enumerate(models_and_neuron_configs.keys()):
+        logger.info(f"***** Compiling {model_name} *****")
         submodel, sub_neuron_config = models_and_neuron_configs[model_name]
         output_file_name = (
             output_file_names[model_name] if output_file_names is not None else Path(model_name + ".neuron")
@@ -298,12 +300,15 @@ def export_models(
         output_path.parent.mkdir(parents=True, exist_ok=True)
 
         try:
+            start_time = time.time()
             neuron_inputs, neuron_outputs = export(
                 model=submodel,
                 config=sub_neuron_config,
                 output=output_path,
                 **compiler_kwargs,
             )
+            compilation_time = time.time() - start_time
+            logger.info(f"[Compilation Time] {np.round(compilation_time, 2)} seconds.")
             outputs.append((neuron_inputs, neuron_outputs))
             # Add neuron specific configs to model components' original config
             if hasattr(submodel, "config"):

optimum/exporters/neuron/model_configs.py

@@ -268,6 +268,13 @@ def outputs(self) -> List[str]:
         return ["sample"]
 
     def generate_dummy_inputs(self, return_tuple: bool = False, **kwargs):
+        # For neuron, we use static shape for compiling the unet. Unlike `optimum`, we use the given `height` and `width` instead of the `sample_size`.
+        if self.height == self.width:
+            self._normalized_config.image_size = self.height
+        else:
+            raise ValueError(
+                "You need to input the same value for `self.height({self.height})` and `self.width({self.width})`."
+            )
         dummy_inputs = super().generate_dummy_inputs(**kwargs)
         dummy_inputs["timestep"] = dummy_inputs["timestep"].float()
         dummy_inputs["encoder_hidden_states"] = dummy_inputs["encoder_hidden_states"][0]

optimum/neuron/modeling_base.py

@@ -15,6 +15,7 @@
 """NeuronBaseModel base classe for inference on neuron devices using the same API as Transformers."""
 
 import logging
+import os
 import shutil
 from contextlib import contextmanager
 from pathlib import Path
@@ -308,6 +309,42 @@ def _from_transformers(
 
         return cls._from_pretrained(save_dir_path, config, model_save_dir=save_dir, neuron_config=neuron_config)
 
+    def push_to_hub(
+        self,
+        save_directory: str,
+        repository_id: str,
+        private: Optional[bool] = None,
+        use_auth_token: Union[bool, str] = True,
+        endpoint: Optional[str] = None,
+    ) -> str:
+        if isinstance(use_auth_token, str):
+            huggingface_token = use_auth_token
+        elif use_auth_token:
+            huggingface_token = HfFolder.get_token()
+        else:
+            raise ValueError("You need to provide `use_auth_token` to be able to push to the hub")
+        api = HfApi(endpoint=endpoint)
+
+        user = api.whoami(huggingface_token)
+        self.git_config_username_and_email(git_email=user["email"], git_user=user["fullname"])
+
+        api.create_repo(
+            token=huggingface_token,
+            repo_id=repository_id,
+            exist_ok=True,
+            private=private,
+        )
+        for path, subdirs, files in os.walk(save_directory):
+            for name in files:
+                local_file_path = os.path.join(path, name)
+                hub_file_path = os.path.relpath(local_file_path, save_directory)
+                api.upload_file(
+                    token=huggingface_token,
+                    repo_id=repository_id,
+                    path_or_fileobj=os.path.join(os.getcwd(), local_file_path),
+                    path_in_repo=hub_file_path,
+                )
+
     def forward(self, *args, **kwargs):
         raise NotImplementedError
 

optimum/neuron/pipelines/diffusers/pipeline_stable_diffusion.py

@@ -158,8 +158,6 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
     def __call__(
         self,
         prompt: Union[str, List[str]] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
         num_inference_steps: int = 50,
         guidance_scale: float = 7.5,
         negative_prompt: Optional[Union[str, List[str]]] = None,
@@ -176,9 +174,9 @@ def __call__(
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
         guidance_rescale: float = 0.0,
     ):
-        # 0. Default height and width to unet
-        height = height or self.unet.config.sample_size * self.vae_scale_factor
-        width = width or self.unet.config.sample_size * self.vae_scale_factor
+        # 0. Height and width to unet (static shapes)
+        height = self.unet.config.neuron["static_height"] * self.vae_scale_factor
+        width = self.unet.config.neuron["static_width"] * self.vae_scale_factor
 
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(