|
|
--- |
|
|
license: apache-2.0 |
|
|
--- |
|
|
|
|
|
Exported with |
|
|
|
|
|
```bash |
|
|
optimum-cli export neuron --model black-forest-labs/FLUX.1-schnell --tensor_parallel_size 8 --batch_size 1 --height 1024 --width 1024 --num_images_per_prompt 1 --sequence_length 256 --torch_dtype bfloat16 flux_schnell_neuron_1024_tp8/ |
|
|
``` |
|
|
|
|
|
Or |
|
|
|
|
|
```python |
|
|
# [Export] |
|
|
import torch |
|
|
from optimum.neuron import NeuronFluxInpaintPipeline |
|
|
|
|
|
if __name__ == "__main__": |
|
|
compiler_args = {"auto_cast": "none"} |
|
|
input_shapes = {"batch_size": 1, "height": 1024, "width": 1024, "sequence_length": 256} |
|
|
|
|
|
pipe = NeuronFluxInpaintPipeline.from_pretrained( |
|
|
"black-forest-labs/FLUX.1-schnell", |
|
|
torch_dtype=torch.bfloat16, |
|
|
export=True, |
|
|
tensor_parallel_size=8, |
|
|
**compiler_args, |
|
|
**input_shapes |
|
|
) |
|
|
|
|
|
# Save locally |
|
|
pipe.save_pretrained("flux_schnell_neuron_1024x1024_tp8/") |
|
|
|
|
|
# Upload to the HuggingFace Hub |
|
|
pipe.push_to_hub( |
|
|
"flux_schnell_neuron_1024x1024_tp8/", repository_id="Jingya/Flux.1-Schnell-1024x1024-neuronx-tp8" # Replace with your HF Hub repo id |
|
|
) |
|
|
``` |
|
|
|