Update app.py

app.py

@@ -2,7 +2,10 @@ import gradio as gr
 import spaces
 import numpy as np
 import random
-from diffusers import DiffusionPipeline
+from diffusers import (
+    DiffusionPipeline, AutoencoderTiny, AutoencoderKL, 
+    AutoPipelineForImage2Image, FluxPipeline, FlowMatchEulerDiscreteScheduler
+)
 import torch
 from PIL import Image
 
@@ -11,14 +14,29 @@ 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 primary diffusion model and assign a smaller VAE for faster real-time previewing
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch_dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae", torch_dtype=torch_dtype).to(device)
+pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype, vae=taef1).to(device)
+
+# Set up for image-to-image pipeline with good VAE and smaller encoder for efficient preview
+pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
+    model_repo_id,
+    vae=good_vae,
+    transformer=pipe.transformer,
+    text_encoder=pipe.text_encoder,
+    tokenizer=pipe.tokenizer,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer_2=pipe.tokenizer_2,
+    torch_dtype=torch_dtype
+)
+
+# Load LoRA weights and set the scale for "hyper-realistic" prompt style
 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
+MAX_SEED = 2**32 - 1
 MAX_IMAGE_SIZE = 1024
 
 # Define styles
@@ -104,6 +122,9 @@ def infer(
 
     return grid_img, seed
 
+# Setup for real-time image generation
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = pipe.flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+
 examples = [
     "A tiny astronaut hatching from an egg on the moon, 4k, planet theme",
     "An anime illustration of a wiener schnitzel --style raw5, 4K",