Update app.py

app.py

@@ -31,6 +31,8 @@ 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_transformer = PeftModel.from_pretrained(repo_large_transformer_path, "jasperai/flash-sd3")
 
 css = '''
 .gradio-container{max-width: 560px !important}
@@ -54,7 +56,7 @@ repo_customs = {
     "Realistic": StableDiffusionXLPipeline.from_pretrained("ehristoforu/Visionix-alpha", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False),
     "Anime": StableDiffusionXLPipeline.from_pretrained("cagliostrolab/animagine-xl-3.1", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False),
     "Pixel": repo_pixel,
-    "Large": StableDiffusion3Pipeline.from_pretrained(repo_large_path, torch_dtype=torch.float16, use_safetensors=True),
+    "Large": StableDiffusion3Pipeline.from_pretrained(repo_large_path, transformer=repo_large_transformer, torch_dtype=torch.float16, use_safetensors=True),
 }
 
 # Functions
@@ -98,8 +100,8 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
         steps_set = 15
         guidance_set = 1.5
     elif model == "Large":   
-        steps_set = 30
-        guidance_set = 5
+        steps_set = 8
+        guidance_set = 1
     else:
         steps_set = 25
         guidance_set = 7