Refactor app.py to optimize model loading and memory management; added AutoencoderTiny and adjusted GPU settings.

app.py

@@ -7,11 +7,8 @@ from gradio_litmodel3d import LitModel3D
 
 import os
 os.environ['SPCONV_ALGO'] = 'native'
-os.system("nvidia-smi")
-os.system("nvcc --version")
 from typing import *
 import torch
-print(f"Torch version {torch.__version__}")
 import numpy as np
 import imageio
 import uuid
@@ -21,15 +18,24 @@ 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, AutoencoderKL
+from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
 from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
 
 llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device)
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
+torch.cuda.empty_cache()
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to("cuda")
-good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = "/tmp/Trellis-demo"
+
+os.makedirs(TMP_DIR, exist_ok=True)
+
 def generate_t2i_prompt(item_name):
     llm_prompt_template = """You are tasked with creating a concise yet highly detailed description of an item to be used for generating an image in a game development pipeline. The image should show the **entire item** with no parts cropped or hidden. The background should always be plain and monocolor, with no focus on it.  
 
@@ -75,7 +81,7 @@ def preprocess_pil_image(image: Image.Image) -> Tuple[str, Image.Image]:
     processed_image.save(f"{TMP_DIR}/{trial_id}.png")
     return trial_id, processed_image
 
-@spaces.GPU
+@spaces.GPU(duration=75)
 def generate_item_image(object_t2i_prompt):
     trial_id = ""
     for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
@@ -103,11 +109,6 @@ def generate_item_image(object_t2i_prompt):
     trial_id, processed_image = preprocess_pil_image(image)
     yield trial_id, processed_image
 
-MAX_SEED = np.iinfo(np.int32).max
-TMP_DIR = "/tmp/Trellis-demo"
-
-os.makedirs(TMP_DIR, exist_ok=True)
-
 def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
     return {
         'gaussian': {