Update app.py

app.py

@@ -3,6 +3,7 @@ from transformers import AutoModel
 from diffusers import FluxPipeline
 import torch
 from huggingface_hub.utils import HfHubHTTPError, RepositoryNotFoundError
+from optimum.neuron import NeuronFluxPipeline
 
 def check_model_access(model_id: str, oauth_token: gr.OAuthToken | None):
     """
@@ -27,7 +28,17 @@ def check_model_access(model_id: str, oauth_token: gr.OAuthToken | None):
     try:
         # 3. The core test: try to load the model's config using the user's token.
         # 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=model_id, torch_dtype=torch.bfloat16, token=oauth_token.token)
+        # pipe = FluxPipeline.from_pretrained(pretrained_model_name_or_path=model_id, torch_dtype=torch.bfloat16, token=oauth_token.token)
+
+        model = NeuronFluxPipeline.from_pretrained(
+            model_id,
+            export=True,
+            tensor_parallel_size=4,
+            token=oauth_token.token,
+            batch_size=1,
+            sequence_length=128,
+        )
+
         
         # 4. If the call succeeds, format a success message
         return f"""