Update app.py

app.py

@@ -30,8 +30,19 @@ DEFAULT_WIDTH = 1024
 
 headers = {"Content-Type": "application/json", "Authorization": f"Bearer {HF_TOKEN}" }
 
-repo_large_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)
-repo_large_transformer_path = SD3Transformer2DModel.from_pretrained(repo_large_path, torch_dtype=torch.float16)
+repo_large_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)
+repo_large_transformer_path = SD3Transformer2DModel.from_pretrained(
+    repo_large_path,
+    subfolder="transformer",
+    torch_dtype=torch.float16,
+)
+
 repo_large_transformer = PeftModel.from_pretrained(repo_large_transformer_path, "jasperai/flash-sd3")
 
 css = '''