Update app.py

app.py

@@ -3,8 +3,18 @@ from diffusers import StableDiffusion3Pipeline
 import gradio as gr
 import os
 import spaces
+from huggingface_hub import snapshot_download
 
-HF_TOKEN = os.environ.get('HF_TOKEN')
+HF_TOKEN = os.getenv("HF_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=HF_TOKEN,
+    )
 
 if torch.cuda.is_available():
     device = "cuda"
@@ -14,7 +24,7 @@ else:
     print("Using CPU")
 
 # Initialize the pipeline and download the model
-pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium", torch_dtype=torch.float16, use_auth_token=HF_TOKEN)
+pipe = StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=torch.float16)
 pipe.to(device)
 
 # Define the image generation function