Update app.py

app.py

@@ -8,17 +8,11 @@ from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlashFlow
 from peft import PeftModel
 import os
 from huggingface_hub import snapshot_download
+auth_token = os.environ.get("TOKEN_FROM_SECRET") or True
 
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 
-model_path = snapshot_download(
-    repo_id="stabilityai/stable-diffusion-3-medium", 
-    revision="refs/pr/26",
-    repo_type="model", 
-    ignore_patterns=["*.md", "*..gitattributes"],
-    local_dir="stable-diffusion-3-medium",
-    token=huggingface_token, # type a new token-id.
-    )
+model_path = "stabilityai/stable-diffusion-3-medium-diffusers"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 IS_SPACE = os.environ.get("SPACE_ID", None) is not None
@@ -27,6 +21,7 @@ transformer = SD3Transformer2DModel.from_pretrained(
     model_path,
     subfolder="transformer",
     torch_dtype=torch.float16,
+    use_auth_token=auth_token
 )
 transformer = PeftModel.from_pretrained(transformer, "jasperai/flash-sd3")
 
@@ -38,7 +33,8 @@ if torch.cuda.is_available():
         transformer=transformer,
         torch_dtype=torch.float16,
         text_encoder_3=None,
-        tokenizer_3=None
+        tokenizer_3=None,
+        use_auth_token=auth_token
     )
 
     pipe = pipe.to(device)
@@ -48,7 +44,8 @@ else:
         transformer=transformer,
         torch_dtype=torch.float16,
         text_encoder_3=None,
-        tokenizer_3=None
+        tokenizer_3=None,
+        use_auth_token=auth_token
     )
     pipe = pipe.to(device)