Upload app.py with huggingface_hub

app.py

@@ -6,13 +6,19 @@ from safetensors.torch import load_file
 
 # Load the model
 pipe = FluxPipeline.from_pretrained(
-    'black-forest-labs/FLUX.1-dev',
+    './data/flux',
     torch_dtype=torch.bfloat16,
     use_safetensors=True
 ).to('cuda')
 
-# Load SRPO weights from https://huggingface.co/tencent/SRPO
-state_dict = load_file("tencent/SRPO/diffusion_pytorch_model.safetensors")
+# Load SRPO weights
+from huggingface_hub import hf_hub_download
+
+srpo_path = hf_hub_download(
+    repo_id="tencent/SRPO",
+    filename="diffusion_pytorch_model.safetensors"
+)
+state_dict = load_file(srpo_path)
 pipe.transformer.load_state_dict(state_dict)
 
 @spaces.GPU(duration=120)