Spaces:
Configuration error
Configuration error
| # Prediction interface for Cog ⚙️ | |
| # https://cog.run/python | |
| from cog import BasePredictor, Input, Path, BaseModel | |
| import os | |
| os.environ['ATTN_BACKEND'] = 'xformers' | |
| import torch | |
| import numpy as np | |
| import imageio | |
| from PIL import Image | |
| from trellis.pipelines import TrellisImageTo3DPipeline | |
| from trellis.utils import render_utils, postprocessing_utils | |
| import logging | |
| MAX_SEED = np.iinfo(np.int32).max | |
| class PredictOutput(BaseModel): | |
| no_background_images: list[Path] | None = None | |
| color_video: Path | None = None | |
| normal_video: Path | None = None | |
| combined_video: Path | None = None | |
| model_file: Path | None = None | |
| gaussian_ply: Path | None = None | |
| class Predictor(BasePredictor): | |
| def setup(self): | |
| """Load the model into memory to make running multiple predictions efficient""" | |
| logging.basicConfig(level=logging.INFO) | |
| self.logger = logging.getLogger(__name__) | |
| self.logger.info("Setting up environment variables...") | |
| os.environ['SPCONV_ALGO'] = 'native' | |
| os.environ['ATTN_BACKEND'] = 'xformers' | |
| self.logger.info("Loading TRELLIS pipeline...") | |
| self.pipeline = TrellisImageTo3DPipeline.from_pretrained("gqk/TRELLIS-image-large-fork") | |
| self.pipeline.cuda() | |
| self.logger.info("Preloading rembg...") | |
| try: | |
| self.pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))) | |
| except Exception as e: | |
| self.logger.warning(f"Rembg preload warning (this is usually fine): {str(e)}") | |
| self.logger.info("Setup complete!") | |
| def predict( | |
| self, | |
| images: list[Path] = Input(description="List of input images to generate 3D asset from"), | |
| seed: int = Input(description="Random seed for generation", default=0), | |
| randomize_seed: bool = Input(description="Randomize seed", default=True), | |
| generate_color: bool = Input(description="Generate color video render", default=True), | |
| generate_normal: bool = Input(description="Generate normal video render", default=False), | |
| generate_model: bool = Input(description="Generate 3D model file (GLB)", default=False), | |
| save_gaussian_ply: bool = Input(description="Save Gaussian point cloud as PLY file", default=False), | |
| return_no_background: bool = Input(description="Return the preprocessed images without background", default=False), | |
| ss_guidance_strength: float = Input( | |
| description="Stage 1: Sparse Structure Generation - Guidance Strength", | |
| default=7.5, | |
| ge=0.0, | |
| le=10.0 | |
| ), | |
| ss_sampling_steps: int = Input( | |
| description="Stage 1: Sparse Structure Generation - Sampling Steps", | |
| default=12, | |
| ge=1, | |
| le=50 | |
| ), | |
| slat_guidance_strength: float = Input( | |
| description="Stage 2: Structured Latent Generation - Guidance Strength", | |
| default=3.0, | |
| ge=0.0, | |
| le=10.0 | |
| ), | |
| slat_sampling_steps: int = Input( | |
| description="Stage 2: Structured Latent Generation - Sampling Steps", | |
| default=12, | |
| ge=1, | |
| le=50 | |
| ), | |
| mesh_simplify: float = Input( | |
| description="GLB Extraction - Mesh Simplification (only used if generate_model=True)", | |
| default=0.95, | |
| ge=0.9, | |
| le=0.98 | |
| ), | |
| texture_size: int = Input( | |
| description="GLB Extraction - Texture Size (only used if generate_model=True)", | |
| default=1024, | |
| ge=512, | |
| le=2048 | |
| ) | |
| ) -> PredictOutput: | |
| """Run a single prediction on the model""" | |
| # Load and process images | |
| self.logger.info("Loading and preprocessing input images...") | |
| input_images = [Image.open(str(image)) for image in images] | |
| processed_images = [self.pipeline.preprocess_image(img) for img in input_images] | |
| # Save the processed images (without background) | |
| no_bg_paths = [] | |
| if return_no_background: | |
| for idx, processed_image in enumerate(processed_images): | |
| no_bg_path = Path(f"output_no_background_{idx}.png") | |
| processed_image.save(str(no_bg_path)) | |
| no_bg_paths.append(no_bg_path) | |
| self.logger.info("Saved images without background") | |
| # Randomize seed if requested | |
| if randomize_seed: | |
| seed = np.random.randint(0, MAX_SEED) | |
| self.logger.info(f"Using randomized seed: {seed}") | |
| else: | |
| self.logger.info(f"Using provided seed: {seed}") | |
| # Generate 3D asset | |
| self.logger.info("Running TRELLIS pipeline...") | |
| if len(processed_images) > 1: | |
| outputs = self.pipeline.run_multi_image( | |
| processed_images, | |
| 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( | |
| processed_images[0], | |
| 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, | |
| } | |
| ) | |
| self.logger.info("TRELLIS pipeline complete!") | |
| self.logger.info(f"Available output formats: {outputs.keys()}") | |
| # Initialize output paths as None | |
| color_path = None | |
| normal_path = None | |
| combined_path = None | |
| model_path = None | |
| gaussian_path = None | |
| # Render videos if requested | |
| if generate_color or generate_normal: | |
| self.logger.info("Starting video rendering...") | |
| if generate_color and generate_normal: | |
| # Generate both videos and combine them side by side | |
| self.logger.info("Generating color video from gaussian output...") | |
| color_renders = render_utils.render_video(outputs['gaussian'][0], num_frames=120) | |
| self.logger.info(f"Available gaussian render types: {list(color_renders.keys())}") | |
| self.logger.info("Generating normal video from mesh output...") | |
| normal_renders = render_utils.render_video(outputs['mesh'][0], num_frames=120) | |
| self.logger.info(f"Available mesh render types: {list(normal_renders.keys())}") | |
| if 'color' in color_renders and 'normal' in normal_renders: | |
| self.logger.info("Combining color and normal videos side by side...") | |
| color_video = color_renders['color'] | |
| normal_video = normal_renders['normal'] | |
| combined_video = [np.concatenate([color_video[i], normal_video[i]], axis=1) for i in range(len(color_video))] | |
| # Save combined video only | |
| combined_path = Path("output_combined.mp4") | |
| imageio.mimsave(str(combined_path), combined_video, fps=15) | |
| self.logger.info("Generated combined video successfully") | |
| else: | |
| self.logger.warning("Missing required render types!") | |
| if 'color' not in color_renders: | |
| self.logger.warning("No color render type found in gaussian output!") | |
| if 'normal' not in normal_renders: | |
| self.logger.warning("No normal render type found in mesh output!") | |
| else: | |
| if generate_color: | |
| self.logger.info("Generating color video from gaussian output...") | |
| color_renders = render_utils.render_video(outputs['gaussian'][0], num_frames=120) | |
| self.logger.info(f"Available gaussian render types: {list(color_renders.keys())}") | |
| if 'color' in color_renders: | |
| color_path = Path("output_color.mp4") | |
| imageio.mimsave(str(color_path), color_renders['color'], fps=15) | |
| self.logger.info("Generated color video successfully") | |
| else: | |
| self.logger.warning("No color render type found in gaussian output!") | |
| if generate_normal: | |
| self.logger.info("Generating normal video from mesh output...") | |
| normal_renders = render_utils.render_video(outputs['mesh'][0], num_frames=120) | |
| self.logger.info(f"Available mesh render types: {list(normal_renders.keys())}") | |
| if 'normal' in normal_renders: | |
| normal_path = Path("output_normal.mp4") | |
| imageio.mimsave(str(normal_path), normal_renders['normal'], fps=15) | |
| self.logger.info("Generated normal video successfully") | |
| else: | |
| self.logger.warning("No normal render type found in mesh output!") | |
| self.logger.info("Video rendering complete!") | |
| # Generate GLB only if requested | |
| if generate_model: | |
| self.logger.info("Generating GLB model...") | |
| glb = postprocessing_utils.to_glb( | |
| outputs['gaussian'][0], | |
| outputs['mesh'][0], | |
| simplify=mesh_simplify, | |
| texture_size=texture_size, | |
| verbose=False | |
| ) | |
| model_path = Path("output.glb") | |
| glb.export(str(model_path)) | |
| self.logger.info("GLB model generation complete!") | |
| # Save Gaussian PLY if requested | |
| if save_gaussian_ply: | |
| self.logger.info("Saving Gaussian point cloud as PLY...") | |
| gaussian_path = Path("output_gaussian.ply") | |
| outputs['gaussian'][0].save_ply(str(gaussian_path)) | |
| self.logger.info("Gaussian PLY file saved successfully!") | |
| self.logger.info("Prediction complete! Returning results...") | |
| return PredictOutput( | |
| no_background_images=no_bg_paths if return_no_background else None, | |
| color_video=color_path if (generate_color and not generate_normal) else None, | |
| normal_video=normal_path if (generate_normal and not generate_color) else None, | |
| combined_video=combined_path if (generate_color and generate_normal) else None, | |
| model_file=model_path, | |
| gaussian_ply=gaussian_path if save_gaussian_ply else None | |
| ) | |