File size: 1,081 Bytes
346dc8c 750584f 346dc8c 4bfbe39 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
---
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
)
```
|