Update image generation pipeline by adding Autoencoder models and random seed for variability

app.py

@@ -18,13 +18,16 @@ from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 from gradio_client import Client
-from diffusers import FluxPipeline
+from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
 from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
+import random
 
 llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = FluxPipeline.from_pretrained("Freepik/flux.1-lite-8B-alpha", torch_dtype=torch.bfloat16).to(device)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device)
+good_vae = AutoencoderKL.from_pretrained("Freepik/flux.1-lite-8B-alpha", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
+pipe = FluxPipeline.from_pretrained("Freepik/flux.1-lite-8B-alpha", torch_dtype=torch.bfloat16, vae=taef1).to(device)
 pipe.vae.enable_tiling()
 pipe.vae.enable_slicing()
 # pipe.enable_sequential_cpu_offload() # offloads modules to CPU on a submodule level (rather than model level)
@@ -85,17 +88,18 @@ def preprocess_pil_image(image: Image.Image) -> Tuple[str, Image.Image]:
 @spaces.GPU(duration=75)
 def generate_item_image(object_t2i_prompt):
     trial_id = ""
+    seed = random.randint(0, MAX_SEED)
     for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
         prompt=object_t2i_prompt,
         guidance_scale=3.5,
-        num_inference_steps=1,
+        num_inference_steps=28,
         width=512,
         height=512,
-        generator=torch.Generator("cpu").manual_seed(0),
+        generator=torch.Generator().manual_seed(seed),
         output_type="pil",
+        good_vae=good_vae,
     ):
         yield trial_id, image
-        print("Yielded image")
 
     trial_id, processed_image = preprocess_pil_image(image)
     yield trial_id, processed_image