Refactor image generation in app.py to streamline processing and enhance performance

app.py

@@ -1,19 +1,11 @@
-import spaces
-from diffusers import FluxPipeline, AutoencoderKL
-from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
-import torch
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-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)
-pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
-
 import gradio as gr
+import spaces
 from gradio_litmodel3d import LitModel3D
 
 import os
 os.environ['SPCONV_ALGO'] = 'native'
 from typing import *
+import torch
 import numpy as np
 import imageio
 import uuid
@@ -23,9 +15,13 @@ 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
 
 llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
+
 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.  
 
@@ -55,32 +51,19 @@ Focus on the item itself, ensuring it is fully described, and specify a plain, w
 
     return object_t2i_prompt
 
-@spaces.GPU(duration=75)
+@spaces.GPU(duration=100)
 def generate_item_image(object_t2i_prompt):
-    trial_id = ""
-    for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-        prompt=object_t2i_prompt,
-        guidance_scale=3.5,
-        num_inference_steps=28,
-        width=1024,
-        height=1024,
-        generator=torch.Generator("cpu").manual_seed(0),
-        output_type="pil",
-        good_vae=good_vae,
-    ):
-        yield trial_id, image
-
-    trial_id, processed_image = preprocess_image(image)
+    image = pipe(prompt=object_t2i_prompt, guidance_scale=3.5, num_inference_steps=28, width=1024, height=1024, generator=torch.Generator("cpu").manual_seed(0), output_type="pil").images[0]
+    trial_id, processed_image = preprocess_pil_image(image)
     return trial_id, processed_image
 
-
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = "/tmp/Trellis-demo"
 
 os.makedirs(TMP_DIR, exist_ok=True)
 
 
-def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
+def preprocess_pil_image(image: Image.Image) -> Tuple[str, Image.Image]:
     """
     Preprocess the input image.
 
@@ -265,7 +248,7 @@ with gr.Blocks(title="Game Items Generator") as demo:
                 for image in os.listdir("assets/example_image")
             ],
             inputs=[image_prompt],
-            fn=preprocess_image,
+            fn=preprocess_pil_image,
             outputs=[trial_id, image_prompt],
             run_on_click=True,
             examples_per_page=64,
@@ -283,7 +266,7 @@ with gr.Blocks(title="Game Items Generator") as demo:
         outputs=[trial_id, image_prompt],
     )
     image_prompt.upload(
-        preprocess_image,
+        preprocess_pil_image,
         inputs=[image_prompt],
         outputs=[trial_id, image_prompt],
     )