Enhance image generation by moving FluxPipeline and AutoencoderKL to GPU; update requirements to include gradio_litmodel3d

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import spaces
-# from gradio_litmodel3d import LitModel3D
+from gradio_litmodel3d import LitModel3D
 
 import os
 os.environ['SPCONV_ALGO'] = 'native'
@@ -20,7 +20,8 @@ from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_imag
 
 llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+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)
 
 def generate_t2i_prompt(item_name):
@@ -71,7 +72,6 @@ def preprocess_pil_image(image: Image.Image) -> Tuple[str, Image.Image]:
 @spaces.GPU
 def generate_item_image(object_t2i_prompt):
     trial_id = ""
-    pipe.to("cuda")
     for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
         prompt=object_t2i_prompt,
         guidance_scale=3.5,
@@ -80,6 +80,7 @@ def generate_item_image(object_t2i_prompt):
         height=1024,
         generator=torch.Generator("cpu").manual_seed(0),
         output_type="pil",
+        good_vae=good_vae,
     ):
         yield trial_id, image
     # img_path = t2i_client.predict(
@@ -93,7 +94,6 @@ def generate_item_image(object_t2i_prompt):
 	# 	api_name="/infer"
     # )[0]
     # image = Image.open(img_path)
-    pipe.to("cpu")
     trial_id, processed_image = preprocess_pil_image(image)
     yield trial_id, processed_image
 
@@ -256,8 +256,8 @@ with gr.Blocks(title="Game Items Generator") as demo:
 
         with gr.Column():
             video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
-            # model_output = LitModel3D(label="Extracted GLB", exposure=20.0, height=300)
-            model_output = gr.Model3D(label="Extracted GLB", height=300)
+            model_output = LitModel3D(label="Extracted GLB", exposure=20.0, height=300)
+            # model_output = gr.Model3D(label="Extracted GLB", height=300)
             download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
             
     trial_id = gr.Textbox(visible=False)

requirements.txt

@@ -21,8 +21,9 @@ pymeshfix==0.17.0
 igraph==0.11.8
 git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
 xformers==0.0.27.post2
-spconv-cu120==2.3.6
+spconv-cu120
 transformers
+gradio_litmodel3d
 https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.0.post2/flash_attn-2.7.0.post2+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
 https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl?download=true
 https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/nvdiffrast-0.3.3-cp310-cp310-linux_x86_64.whl?download=true