Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import spaces
 import numpy as np
 import random
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoencoderTiny
 import torch
 from PIL import Image
 
@@ -11,11 +11,26 @@ model_repo_id = "stabilityai/stable-diffusion-3.5-large-turbo"
 
 torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+# Load Tiny Autoencoder and optimize its decoder layers with torch.compile
+taesd3 = (
+    AutoencoderTiny.from_pretrained("madebyollin/taesd3", torch_dtype=torch.float16)
+    .half()
+    .eval()
+    .requires_grad_(False)
+    .to(device)
+)
+taesd3.decoder.layers = torch.compile(
+    taesd3.decoder.layers,
+    fullgraph=True,
+    dynamic=False,
+    mode="max-autotune-no-cudagraphs",
+)
+
+# Load main Stable Diffusion pipeline with Tiny Autoencoder
+pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype, vae=taesd3).to(device)
 
 pipe.load_lora_weights("prithivMLmods/SD3.5-Large-Turbo-HyperRealistic-LoRA", weight_name="SD3.5-4Step-Large-Turbo-HyperRealistic-LoRA.safetensors")
-trigger_word = "hyper realistic"  # Specify trigger word for LoRA
+trigger_word = "hyper realistic"
 pipe.fuse_lora(lora_scale=1.0)
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -231,4 +246,4 @@ with gr.Blocks(css=css, theme="prithivMLmods/Minecraft-Theme") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()