Update app.py

app.py

@@ -7,12 +7,24 @@ import torch
 from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlashFlowMatchEulerDiscreteScheduler
 from peft import PeftModel
 import os
+from huggingface_hub import snapshot_download
+
+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.
+    )
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 IS_SPACE = os.environ.get("SPACE_ID", None) is not None
 
 transformer = SD3Transformer2DModel.from_pretrained(
-    "stabilityai/stable-diffusion-3-medium-diffusers",
+    model_path,
     subfolder="transformer",
     torch_dtype=torch.float16,
 )
@@ -22,7 +34,7 @@ transformer = PeftModel.from_pretrained(transformer, "jasperai/flash-sd3")
 if torch.cuda.is_available():
     torch.cuda.max_memory_allocated(device=device)
     pipe = StableDiffusion3Pipeline.from_pretrained(
-        "stabilityai/stable-diffusion-3-medium-diffusers",
+        model_path,
         transformer=transformer,
         torch_dtype=torch.float16,
         text_encoder_3=None,
@@ -32,7 +44,7 @@ if torch.cuda.is_available():
     pipe = pipe.to(device)
 else:
     pipe = StableDiffusion3Pipeline.from_pretrained(
-        "stabilityai/stable-diffusion-3-medium-diffusers",
+        model_path,
         transformer=transformer,
         torch_dtype=torch.float16,
         text_encoder_3=None,
@@ -42,7 +54,7 @@ else:
 
 
 pipe.scheduler = FlashFlowMatchEulerDiscreteScheduler.from_pretrained(
-  "stabilityai/stable-diffusion-3-medium-diffusers",
+  model_path,
   subfolder="scheduler",
 )