Update Gradio app with multiple files

models.py

@@ -1,5 +1,6 @@
 import spaces
 import torch
+import os
 from diffusers import DiffusionPipeline
 
 MODEL_ID = 'black-forest-labs/FLUX.1-dev'
@@ -7,7 +8,9 @@ MODEL_ID = 'black-forest-labs/FLUX.1-dev'
 # Compile the model ahead-of-time for optimal performance
 @spaces.GPU(duration=1500)  # Maximum duration for compilation
 def compile_transformer():
-    pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
+    # Load model with HF token if available
+    token = os.getenv('HF_TOKEN')
+    pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, token=token)
     pipe.to('cuda')
     
     # Capture example inputs for AoT compilation
@@ -26,7 +29,9 @@ def compile_transformer():
 
 # Load and compile the model
 def load_model():
-    pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
+    # Load model with HF token if available
+    token = os.getenv('HF_TOKEN')
+    pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, token=token)
     pipe.to('cuda')
     
     # Apply compiled transformer