Spaces:

optimum
/

neuron-export

Paused

badaoui HF Staff commited on Aug 25

Commit

ba399b5

verified ·

1 Parent(s): 203503a

typo

Files changed (1) hide show

app.py CHANGED Viewed

@@ -138,7 +138,7 @@ def neuron_export(model_id: str, model_type: str, task_or_pipeline: str,
         # This will fail if the user doesn't have access to the gated repo's files.
         # pipe = FluxPipeline.from_pretrained(pretrained_model_name_or_path="black-forest-labs/FLUX.1-Krea-dev", torch_dtype=torch.bfloat16, token=oauth_token.token)
-        model = model_class.from_pretrained(
             "black-forest-labs/FLUX.1-Krea-dev",
             export=True,
             tensor_parallel_size=4,

         # This will fail if the user doesn't have access to the gated repo's files.
         # pipe = FluxPipeline.from_pretrained(pretrained_model_name_or_path="black-forest-labs/FLUX.1-Krea-dev", torch_dtype=torch.bfloat16, token=oauth_token.token)
+        model = NeuronFluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-Krea-dev",
             export=True,
             tensor_parallel_size=4,