Flux.1-Schnell-1024x1024-neuronx-tp8

Exported with

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

# [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
    )