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, AutoencoderTiny
+from diffusers import DiffusionPipeline, AutoencoderTiny, StableDiffusion3Pipeline, SD3Transformer2DModel, FlashFlowMatchEulerDiscreteScheduler
 import torch
 from PIL import Image
 
@@ -10,15 +10,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "stabilityai/stable-diffusion-3.5-large-turbo"
 
 torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
-# 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)
-)
+# Load Tiny Autoencoder
+taesd3 = AutoencoderTiny.from_pretrained("madebyollin/taesd3", torch_dtype=torch.float16).to(device)
 taesd3.decoder.layers = torch.compile(
     taesd3.decoder.layers,
     fullgraph=True,
@@ -26,12 +19,16 @@ taesd3.decoder.layers = torch.compile(
     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)
+# Load main Stable Diffusion pipeline
+pipe = StableDiffusion3Pipeline.from_pretrained(
+    model_path,
+    transformer=SD3Transformer2DModel.from_pretrained(model_path, torch_dtype=torch.float16),
+    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"
-pipe.fuse_lora(lora_scale=1.0)
+pipe.scheduler = FlashFlowMatchEulerDiscreteScheduler.from_pretrained(model_path, subfolder="scheduler")
+pipe.set_progress_bar_config(disable=True)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -221,29 +218,14 @@ with gr.Blocks(css=css, theme="prithivMLmods/Minecraft-Theme") as demo:
                     value=4,
                 )
 
-        gr.Examples(examples=examples, 
-                    inputs=[prompt], 
-                    outputs=[result, seed], 
-                    fn=infer, 
-                    cache_examples=False)
+        gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            style_selection,
-            grid_size_selection,
-        ],
+        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, style_selection, grid_size_selection],
         outputs=[result, seed],
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()