Views
No views yet
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/1# [Export]
2import torch
3from optimum.neuron import NeuronFluxInpaintPipeline
4
5if __name__ == "__main__":
6 compiler_args = {"auto_cast": "none"}
7 input_shapes = {"batch_size": 1, "height": 1024, "width": 1024, "sequence_length": 256}
8
9 pipe = NeuronFluxInpaintPipeline.from_pretrained(
10 "black-forest-labs/FLUX.1-schnell",
11 torch_dtype=torch.bfloat16,
12 export=True,
13 tensor_parallel_size=8,
14 **compiler_args,
15 **input_shapes
16 )
17
18 # Save locally
19 pipe.save_pretrained("flux_schnell_neuron_1024x1024_tp8/")
20
21 # Upload to the HuggingFace Hub
22 pipe.push_to_hub(
23 "flux_schnell_neuron_1024x1024_tp8/", repository_id="Jingya/Flux.1-Schnell-1024x1024-neuronx-tp8" # Replace with your HF Hub repo id
24 )