Update app.py

app.py

@@ -1,244 +1,284 @@
 import gradio as gr
-import spaces
-from gradio_litmodel3d import LitModel3D
-import os
-import shutil
 import torch
 import numpy as np
+import os
+import shutil
 import imageio
-from easydict import EasyDict as edict
 from PIL import Image
-from trellis.pipelines import TrellisImageTo3DPipeline
-from trellis.representations import Gaussian, MeshExtractResult
-from trellis.utils import render_utils, postprocessing_utils
 
-# Add this before your preprocessing functions
-pipeline = TrellisImageTo3DPipeline()
+# Ensure imports are available
+try:
+    from trellis.pipelines import TrellisImageTo3DPipeline
+    from trellis.representations import Gaussian, MeshExtractResult
+    from trellis.utils import render_utils, postprocessing_utils
+    from easydict import EasyDict as edict
+except ImportError as e:
+    print(f"Error importing required libraries: {e}")
+    print("Please install the following libraries:")
+    print("- trellis-ai")
+    print("- easydict")
+    TrellisImageTo3DPipeline = None
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(TMP_DIR, exist_ok=True)
 
-# Session management
-def start_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(user_dir, exist_ok=True)
-
-def end_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shutil.rmtree(user_dir)
-
-# Preprocessing functions
-def preprocess_image(image: Image.Image) -> Image.Image:
-    return pipeline.preprocess_image(image)
-
-def preprocess_images(images: list) -> list:
-    images = [image[0] for image in images]
-    return [pipeline.preprocess_image(image) for image in images]
-
-# Utility functions
-def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
-    return {
-        'gaussian': {
-            **gs.init_params,
-            '_xyz': gs._xyz.cpu().numpy(),
-            '_features_dc': gs._features_dc.cpu().numpy(),
-            '_scaling': gs._scaling.cpu().numpy(),
-            '_rotation': gs._rotation.cpu().numpy(),
-            '_opacity': gs._opacity.cpu().numpy(),
-        },
-        'mesh': {
-            'vertices': mesh.vertices.cpu().numpy(),
-            'faces': mesh.faces.cpu().numpy(),
-        },
-    }
-
-def unpack_state(state: dict) -> tuple:
-    gs = Gaussian(
-        aabb=state['gaussian']['aabb'],
-        sh_degree=state['gaussian']['sh_degree'],
-        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
-        scaling_bias=state['gaussian']['scaling_bias'],
-        opacity_bias=state['gaussian']['opacity_bias'],
-        scaling_activation=state['gaussian']['scaling_activation'],
-    )
-    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
-    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
-    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
-    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
-    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
-
-    mesh = edict(
-        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
-        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
-    )
-
-    return gs, mesh
-
-def get_seed(randomize_seed: bool, seed: int) -> int:
-    return np.random.randint(0, MAX_SEED) if randomize_seed else seed
-
-# Core functions
-@spaces.GPU
-def image_to_3d(
-    image: Image.Image,
-    multiimages: list,
-    is_multiimage: bool,
-    seed: int,
-    ss_guidance_strength: float,
-    ss_sampling_steps: int,
-    slat_guidance_strength: float,
-    slat_sampling_steps: int,
-    multiimage_algo: str,
-    req: gr.Request,
-) -> tuple:
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    if not is_multiimage:
-        outputs = pipeline.run(
-            image,
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            preprocess_image=False,
-            sparse_structure_sampler_params={
-                "steps": ss_sampling_steps,
-                "cfg_strength": ss_guidance_strength,
-            },
-            slat_sampler_params={
-                "steps": slat_sampling_steps,
-                "cfg_strength": slat_guidance_strength,
-            },
+class ImageTo3DConverter:
+    def __init__(self):
+        # Initialize the pipeline with error handling
+        try:
+            self.pipeline = TrellisImageTo3DPipeline()
+        except Exception as e:
+            print(f"Failed to initialize pipeline: {e}")
+            self.pipeline = None
+
+    def validate_input(self, image, is_multiimage):
+        """Validate input images before processing"""
+        if not self.pipeline:
+            raise ValueError("Pipeline not initialized. Check library installation.")
+        
+        if is_multiimage:
+            # Handle multi-image input
+            if not image or len(image) == 0:
+                raise ValueError("No images provided for multi-image processing")
+            # Ensure images are PIL Image objects
+            valid_images = [img[0] if isinstance(img, list) else img for img in image]
+            return valid_images
+        else:
+            # Handle single image input
+            if image is None:
+                raise ValueError("No image provided")
+            return image
+
+    def preprocess_image(self, image):
+        """Safely preprocess a single image"""
+        try:
+            return self.pipeline.preprocess_image(image)
+        except Exception as e:
+            print(f"Image preprocessing error: {e}")
+            return image
+
+    def process_image(self, 
+                      image, 
+                      multiimages, 
+                      is_multiimage, 
+                      seed, 
+                      ss_guidance_strength, 
+                      ss_sampling_steps,
+                      slat_guidance_strength, 
+                      slat_sampling_steps,
+                      multiimage_algo):
+        """Main image to 3D conversion method"""
+        # Validate and preprocess input
+        try:
+            processed_input = self.validate_input(image if not is_multiimage else multiimages, is_multiimage)
+        except ValueError as e:
+            print(f"Input validation error: {e}")
+            return None, None
+
+        # Determine processing method based on input type
+        try:
+            if not is_multiimage:
+                outputs = self.pipeline.run(
+                    processed_input,
+                    seed=seed,
+                    formats=["gaussian", "mesh"],
+                    preprocess_image=False,
+                    sparse_structure_sampler_params={
+                        "steps": ss_sampling_steps,
+                        "cfg_strength": ss_guidance_strength,
+                    },
+                    slat_sampler_params={
+                        "steps": slat_sampling_steps,
+                        "cfg_strength": slat_guidance_strength,
+                    },
+                )
+            else:
+                outputs = self.pipeline.run_multi_image(
+                    processed_input,
+                    seed=seed,
+                    formats=["gaussian", "mesh"],
+                    preprocess_image=False,
+                    sparse_structure_sampler_params={
+                        "steps": ss_sampling_steps,
+                        "cfg_strength": ss_guidance_strength,
+                    },
+                    slat_sampler_params={
+                        "steps": slat_sampling_steps,
+                        "cfg_strength": slat_guidance_strength,
+                    },
+                    mode=multiimage_algo,
+                )
+        except Exception as e:
+            print(f"3D conversion error: {e}")
+            return None, None
+
+        # Generate video
+        try:
+            video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+            video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+            video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+            
+            # Save video
+            user_dir = os.path.join(TMP_DIR, 'temp_session')
+            os.makedirs(user_dir, exist_ok=True)
+            video_path = os.path.join(user_dir, 'sample.mp4')
+            imageio.mimsave(video_path, video, fps=15)
+
+            # Pack and return state
+            state = {
+                'gaussian': {
+                    **outputs['gaussian'][0].init_params,
+                    '_xyz': outputs['gaussian'][0]._xyz.cpu().numpy(),
+                    '_features_dc': outputs['gaussian'][0]._features_dc.cpu().numpy(),
+                    '_scaling': outputs['gaussian'][0]._scaling.cpu().numpy(),
+                    '_rotation': outputs['gaussian'][0]._rotation.cpu().numpy(),
+                    '_opacity': outputs['gaussian'][0]._opacity.cpu().numpy(),
+                },
+                'mesh': {
+                    'vertices': outputs['mesh'][0].vertices.cpu().numpy(),
+                    'faces': outputs['mesh'][0].faces.cpu().numpy(),
+                },
+            }
+
+            return state, video_path
+
+        except Exception as e:
+            print(f"Video generation error: {e}")
+            return None, None
+
+    def extract_glb(self, state, mesh_simplify=0.95, texture_size=1024):
+        """Extract GLB from the processed state"""
+        try:
+            # Reconstruct Gaussian and Mesh from state
+            gs = Gaussian(
+                aabb=state['gaussian']['aabb'],
+                sh_degree=state['gaussian']['sh_degree'],
+                mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
+                scaling_bias=state['gaussian'].get('scaling_bias', 0.1),
+                opacity_bias=state['gaussian'].get('opacity_bias', 0.1),
+                scaling_activation=state['gaussian'].get('scaling_activation', 'softplus'),
+            )
+            gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
+            gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
+            gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
+            gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
+            gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
+
+            mesh = edict(
+                vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
+                faces=torch.tensor(state['mesh']['faces'], device='cuda'),
+            )
+
+            # Convert mesh
+            mesh.vertices, mesh.faces = postprocessing_utils.remesh_to_quads(
+                vertices=mesh.vertices.cpu().numpy(),
+                faces=mesh.faces.cpu().numpy(),
+                simplify=mesh_simplify
+            )
+
+            # Generate GLB
+            glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
+            
+            # Save GLB
+            user_dir = os.path.join(TMP_DIR, 'temp_session')
+            os.makedirs(user_dir, exist_ok=True)
+            glb_path = os.path.join(user_dir, 'sample.glb')
+            glb.export(glb_path)
+
+            return glb_path
+
+        except Exception as e:
+            print(f"GLB extraction error: {e}")
+            return None
+
+# Gradio Interface Setup
+def create_gradio_interface():
+    converter = ImageTo3DConverter()
+
+    with gr.Blocks() as demo:
+        # Input components
+        with gr.Row():
+            with gr.Column():
+                with gr.Tabs() as input_tabs:
+                    with gr.Tab("Single Image"):
+                        single_image = gr.Image(label="Single Image Input")
+                    with gr.Tab("Multiple Images"):
+                        multi_images = gr.Gallery(label="Multiple Image Input")
+                
+                # Generation settings
+                with gr.Accordion("Generation Settings"):
+                    seed = gr.Slider(0, MAX_SEED, label="Seed", value=0)
+                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                    
+                    with gr.Row():
+                        ss_guidance = gr.Slider(0, 10, label="Sparse Guidance Strength", value=7.5)
+                        ss_steps = gr.Slider(1, 50, label="Sparse Sampling Steps", value=12)
+                    
+                    with gr.Row():
+                        slat_guidance = gr.Slider(0, 10, label="Latent Guidance Strength", value=3.0)
+                        slat_steps = gr.Slider(1, 50, label="Latent Sampling Steps", value=12)
+                    
+                    multi_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="stochastic")
+
+                # Buttons
+                generate_btn = gr.Button("Generate 3D Model")
+                
+                # GLB Extraction
+                with gr.Accordion("GLB Extraction"):
+                    mesh_simplify = gr.Slider(0.9, 0.98, label="Mesh Simplify", value=0.95)
+                    texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024)
+                    extract_glb_btn = gr.Button("Extract GLB")
+
+        # Output components
+        with gr.Column():
+            video_output = gr.Video(label="Generated 3D Asset Preview")
+            glb_output = gr.File(label="Extracted GLB")
+
+        # Event handlers
+        def generate_3d(image, multi_image, seed, ss_guidance, ss_steps, 
+                        slat_guidance, slat_steps, multi_algo):
+            # Determine if it's multi-image mode
+            is_multi = isinstance(multi_image, list) and len(multi_image) > 0
+            
+            # Randomize seed if selected
+            if randomize_seed:
+                seed = np.random.randint(0, MAX_SEED)
+            
+            # Process image
+            state, video = converter.process_image(
+                image, multi_image, is_multi, seed, 
+                ss_guidance, ss_steps, 
+                slat_guidance, slat_steps, 
+                multi_algo
+            )
+            
+            return video if video else None
+
+        def extract_glb(state, simplify, texture_size):
+            if state is None:
+                return None
+            glb_path = converter.extract_glb(state, simplify, texture_size)
+            return glb_path
+
+        # Connect event handlers
+        generate_btn.click(
+            generate_3d, 
+            inputs=[single_image, multi_images, seed, ss_guidance, ss_steps, 
+                    slat_guidance, slat_steps, multi_algo],
+            outputs=[video_output]
         )
-    else:
-        outputs = pipeline.run_multi_image(
-            [image[0] for image in multiimages],
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            preprocess_image=False,
-            sparse_structure_sampler_params={
-                "steps": ss_sampling_steps,
-                "cfg_strength": ss_guidance_strength,
-            },
-            slat_sampler_params={
-                "steps": slat_sampling_steps,
-                "cfg_strength": slat_guidance_strength,
-            },
-            mode=multiimage_algo,
+
+        extract_glb_btn.click(
+            extract_glb,
+            inputs=[state, mesh_simplify, texture_size],
+            outputs=[glb_output]
         )
 
-    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
-    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
-    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-    video_path = os.path.join(user_dir, 'sample.mp4')
-    imageio.mimsave(video_path, video, fps=15)
-    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
-    torch.cuda.empty_cache()
-    return state, video_path
-
-@spaces.GPU(duration=90)
-def extract_glb(
-    state: dict,
-    mesh_simplify: float,
-    texture_size: int,
-    req: gr.Request,
-) -> tuple:
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    gs, mesh = unpack_state(state)
-
-    # Convert the mesh to polygonal surfaces (quads)
-    mesh.vertices, mesh.faces = postprocessing_utils.remesh_to_quads(
-        vertices=mesh.vertices.cpu().numpy(),
-        faces=mesh.faces.cpu().numpy(),
-        simplify=mesh_simplify
-    )
-
-    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-    glb_path = os.path.join(user_dir, 'sample.glb')
-    glb.export(glb_path)
-    torch.cuda.empty_cache()
-    return glb_path, glb_path
-
-@spaces.GPU
-def extract_gaussian(state: dict, req: gr.Request) -> tuple:
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    gs, _ = unpack_state(state)
-    gaussian_path = os.path.join(user_dir, 'sample.ply')
-    gs.save_ply(gaussian_path)
-    torch.cuda.empty_cache()
-    return gaussian_path, gaussian_path
-
-# Gradio UI setup
-with gr.Blocks(theme=gr.themes.Default(), delete_cache=(600, 600)) as demo:
-    with gr.Row():
-        with gr.Column():
-            with gr.Tabs() as input_tabs:
-                with gr.Tab(label="Single Image", id=0) as single_image_input_tab:
-                    image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300)
-                with gr.Tab(label="Multiple Images", id=1) as multiimage_input_tab:
-                    multiimage_prompt = gr.Gallery(label="Image Prompt", format="png", type="pil", height=300, columns=3)
-
-            with gr.Accordion(label="Generation Settings", open=False):
-                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
-                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                with gr.Row():
-                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Sparse Guidance Strength", value=7.5, step=0.1)
-                    ss_sampling_steps = gr.Slider(1, 50, label="Sparse Sampling Steps", value=12, step=1)
-                with gr.Row():
-                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Latent Guidance Strength", value=3.0, step=0.1)
-                    slat_sampling_steps = gr.Slider(1, 50, label="Latent Sampling Steps", value=12, step=1)
-                multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="stochastic")
-
-            generate_btn = gr.Button("Generate", variant="primary")
-
-            with gr.Accordion(label="GLB Extraction Settings", open=False):
-                mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
-                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
-
-            with gr.Row():
-                extract_glb_btn = gr.Button("Extract GLB", interactive=False)
-                extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
+    return demo
 
-        with gr.Column():
-            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
-            model_output = LitModel3D(label="Extracted GLB/Gaussian", exposure=10.0, height=300)
-
-            with gr.Row():
-                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
-                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)
-
-    is_multiimage = gr.State(False)
-    output_buf = gr.State()
-
-    # Handlers
-    demo.load(start_session)
-    demo.unload(end_session)
-
-    single_image_input_tab.select(lambda: False, outputs=[is_multiimage])
-    multiimage_input_tab.select(lambda: True, outputs=[is_multiimage])
-
-    image_prompt.upload(preprocess_image, inputs=[image_prompt], outputs=[image_prompt])
-    multiimage_prompt.upload(preprocess_images, inputs=[multiimage_prompt], outputs=[multiimage_prompt])
-
-    generate_btn.click(get_seed, inputs=[randomize_seed, seed], outputs=[seed]).then(
-        image_to_3d,
-        inputs=[image_prompt, multiimage_prompt, is_multiimage, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo],
-        outputs=[output_buf, video_output],
-    ).then(
-        lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
-        outputs=[extract_glb_btn, extract_gs_btn],
-    )
-
-    extract_glb_btn.click(
-        extract_glb,
-        inputs=[output_buf, mesh_simplify, texture_size],
-        outputs=[model_output, download_glb],
-    )
-
-    extract_gs_btn.click(
-        extract_gaussian,
-        inputs=[output_buf],
-        outputs=[model_output, download_gs],
-    )
-
-# Launch the Gradio demo for Hugging Face Spaces
-demo.launch()
+# Launch the interface
+if __name__ == "__main__":
+    interface = create_gradio_interface()
+    interface.launch()