import gradio as gr import spaces import os import shutil os.environ['SPCONV_ALGO'] = 'native' from typing import * import torch import numpy as np 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 import requests import base64 import io import tempfile 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) NODE_SERVER_UPLOAD_URL = "https://viverse-backend.onrender.com/api/upload-rigged-model" # Funciones auxiliares 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) def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]: images = [image[0] for image in images] processed_images = [pipeline.preprocess_image(image) for image in images] return processed_images 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[Gaussian, edict]: 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 @spaces.GPU def image_to_3d( multiimages: List[Tuple[Image.Image, str]], seed: int, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int, multiimage_algo: Literal["multidiffusion", "stochastic"], req: gr.Request, ) -> Tuple[dict, str]: user_dir = os.path.join(TMP_DIR, str(req.session_hash)) 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, ) 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[str, str]: user_dir = os.path.join(TMP_DIR, str(req.session_hash)) gs, mesh = unpack_state(state) 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(duration=180) def generate_model_from_images_and_upload( image_inputs: List[Dict[str, Any]], input_type: str, seed_val: int, ss_guidance_strength_val: float, ss_sampling_steps_val: int, slat_guidance_strength_val: float, slat_sampling_steps_val: int, multiimage_algo_val: str, mesh_simplify_val: float, texture_size_val: int, model_description: str, req: gr.Request ) -> str: user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) # --- DEBUG LOGS --- print(f"Python DEBUG: Raw image_inputs (as received by function): {image_inputs}") print(f"Python DEBUG: Type of image_inputs: {type(image_inputs)}") if isinstance(image_inputs, list): print(f"Python DEBUG: Length of image_inputs list: {len(image_inputs)}") if len(image_inputs) > 0 and isinstance(image_inputs[0], dict): print(f"Python DEBUG: First element of image_inputs (should be a dict): {image_inputs[0]}") print(f"Python DEBUG: Type of first element: {type(image_inputs[0])}") print(f"Python DEBUG: Received input_type from Node.js: '{input_type}'") # Should always be 'url' now # --- END DEBUG LOGS --- pil_images = [] image_basenames_for_prompt = [] for i, file_data_obj in enumerate(image_inputs): # file_data_obj is one dict from the list img_to_open_path = None current_image_name = file_data_obj.get('name', f"image_{i}.png") print(f"Python DEBUG: Processing item {i}: {file_data_obj}, current_image_name: {current_image_name}") # For URLs (which is now always the case from Node.js), # Gradio should have downloaded the image and put its local path in file_data_obj.get('path') img_to_open_path = file_data_obj.get('path') if not img_to_open_path: print(f"Error: 'path' was missing in item {i}: {file_data_obj}. Skipping.") continue print(f"Python INFO: Using Gradio-provided path for '{current_image_name}': {img_to_open_path}") # Now, process the image using img_to_open_path try: print(f"Python INFO: Opening image from path: {img_to_open_path} (intended name for prompt: {current_image_name})") img = Image.open(img_to_open_path) image_basenames_for_prompt.append(os.path.splitext(current_image_name)[0] or f"image_{i}") if img.mode == 'RGBA' or img.mode == 'P': print(f"Converting image '{current_image_name}' from {img.mode} to RGB") img = img.convert('RGB') processed_img = pipeline.preprocess_image(img) pil_images.append(processed_img) print(f"Image '{current_image_name}' (item {i+1}) processed successfully and added to list.") except Exception as e_img_proc: print(f"Error opening or processing image at '{img_to_open_path}' (item {i}, name: '{current_image_name}'): {e_img_proc}") import traceback traceback.print_exc() # Continue to next image if one fails # No finally block needed here anymore for deleting temp base64 files if not pil_images: print("Error: No images could be processed from the input. Aborting generation.") raise gr.Error("Failed to process any input images.") print(f"Python INFO: Total images processed for pipeline: {len(pil_images)}") effective_model_description = model_description if not effective_model_description and image_basenames_for_prompt: effective_model_description = "_prompted_by_" + "_and_".join(image_basenames_for_prompt) effective_model_description = effective_model_description[:100] # Keep it reasonably short elif not effective_model_description: effective_model_description = "ImageGenModel" print(f"Python INFO: Using model_description for upload: {effective_model_description}") # Generate 3D model using the Trellis image pipeline try: print(f"Python INFO: Calling internal image_to_3d with {len(pil_images)} images.") # The image_to_3d function expects a list of tuples (PIL.Image, str_filename_or_label) # We have processed_img in pil_images which are already PIL.Image objects after pipeline.preprocess_image # We can use the current_image_name (or derived basenames) as the string part if needed by image_to_3d, # but Trellis's run_multi_image takes a list of PIL images directly. # Let's adapt multiimages for image_to_3d to be List[Tuple[Image.Image, str]] multiimages_for_pipeline = [] for idx, p_img in enumerate(pil_images): # Create a simple label for each image for the tuple structure label = image_basenames_for_prompt[idx] if idx < len(image_basenames_for_prompt) else f"image_{idx}" multiimages_for_pipeline.append((p_img, label)) # p_img here is already the *processed* image tensor. # image_to_3d will take image[0] from this list. # This might need adjustment if image_to_3d expects raw PIL Images. # Re-checking Trellis: pipeline.run_multi_image takes List[Image.Image] # and preprocesses them internally if preprocess_image=True (default). # Since we pre-process above, we should pass preprocess_image=False if run_multi_image allows. # The `image_to_3d` in this file is a wrapper for run_multi_image. # It passes `preprocess_image=False`. # So, `pil_images` containing already processed images is what `image_to_3d` expects for `[image[0] for image in multiimages]` state, _ = image_to_3d( multiimages=[(img, name) for img, name in zip(pil_images, image_basenames_for_prompt)], # Pass list of (processed_PIL_image, name_str) seed=seed_val, ss_guidance_strength=ss_guidance_strength_val, ss_sampling_steps=ss_sampling_steps_val, slat_guidance_strength=slat_guidance_strength_val, slat_sampling_steps=slat_sampling_steps_val, multiimage_algo=multiimage_algo_val, req=req ) if state is None: print("Error: Internal image_to_3d returned None state!") raise ValueError("Internal image_to_3d failed to return state") print(f"Python INFO: Internal image_to_3d completed. State type: {type(state)}") print("Python INFO: Calling internal extract_glb...") glb_path, _ = extract_glb( state, mesh_simplify_val, texture_size_val, req ) if glb_path is None or not os.path.isfile(glb_path): print(f"Error: Internal extract_glb returned None or invalid path: {glb_path}") raise FileNotFoundError(f"Generated GLB file not found at {glb_path}") print(f"Python INFO: Internal extract_glb completed. GLB path: {glb_path}") print(f"Python INFO: Uploading GLB from {glb_path} to {NODE_SERVER_UPLOAD_URL}") persistent_url = None with open(glb_path, "rb") as f: files = {"modelFile": (os.path.basename(glb_path), f, "model/gltf-binary")} payload = { "clientType": "imagen", "modelStage": "imagen_mesh", "prompt": effective_model_description # Use the description here } print(f"Python INFO: Upload payload: {payload}") response = requests.post(NODE_SERVER_UPLOAD_URL, files=files, data=payload) response.raise_for_status() # Raise an exception for bad status codes result = response.json() persistent_url = result.get("persistentUrl") if not persistent_url: print(f"Error: No persistent URL in Node.js server response: {result}") raise ValueError("Upload successful, but no persistent URL returned") print(f"Python INFO: Successfully uploaded to Node server. Persistent URL: {persistent_url}") return persistent_url except Exception as e: print(f"ERROR in Image-to-3D pipeline: {e}") import traceback traceback.print_exc() raise gr.Error(f"Image-to-3D pipeline failed: {e}") # Interfaz Gradio with gr.Blocks(delete_cache=(600, 600)) as demo: gr.Markdown(""" # UTPL - Conversión de Multiples Imágenes a objetos 3D usando IA ### Tesis: *"Objetos tridimensionales creados por IA: Innovación en entornos virtuales"* **Autor:** Carlos Vargas **Base técnica:** Adaptación de [TRELLIS](https://trellis3d.github.io/) (herramienta de código abierto para generación 3D) **Propósito educativo:** Demostraciones académicas e Investigación en modelado 3D automático """) with gr.Row(): with gr.Column(): with gr.Tabs() as input_tabs: 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) gr.Markdown("Stage 1: Sparse Structure Generation") with gr.Row(): ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1) ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1) gr.Markdown("Stage 2: Structured Latent Generation") with gr.Row(): slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1) slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1) multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="stochastic") generate_btn = gr.Button("Generate") 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) extract_glb_btn = gr.Button("Extract GLB", interactive=False) with gr.Column(): video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300) model_output = gr.Model3D(label="Extracted GLB", height=300) download_glb = gr.DownloadButton(label="Download GLB", interactive=False) output_buf = gr.State() # Manejadores demo.load(start_session) demo.unload(end_session) 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=[multiimage_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo], outputs=[output_buf, video_output], ).then( lambda: gr.Button(interactive=True), outputs=[extract_glb_btn], ) video_output.clear( lambda: gr.Button(interactive=False), outputs=[extract_glb_btn], ) extract_glb_btn.click( extract_glb, inputs=[output_buf, mesh_simplify, texture_size], outputs=[model_output, download_glb], ).then( lambda: gr.Button(interactive=True), outputs=[download_glb], ) model_output.clear( lambda: gr.Button(interactive=False), outputs=[download_glb], ) # --- Add this section to explicitly register the API function for image to 3D --- # Use gr.JSON for the API input, as it's designed for arbitrary JSON data. api_image_inputs_json = gr.JSON(value=[]) # <--- CHANGED to gr.JSON and new variable name api_input_type_state = gr.State(value="url") api_model_description_state = gr.State(value="ImagenModel") with gr.Row(visible=False): # Hide this row in the UI api_image_gen_trigger_btn = gr.Button("API Image-to-3D Trigger") api_image_gen_output_url = gr.Textbox(label="Generated Model URL (API)", visible=False) api_image_gen_trigger_btn.click( generate_model_from_images_and_upload, inputs=[ api_image_inputs_json, # <--- CHANGED to use the gr.JSON component's variable name api_input_type_state, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo, mesh_simplify, texture_size, api_model_description_state, ], outputs=[api_image_gen_output_url], api_name="generate_model_from_images_and_upload" ) # --- End API registration section --- # Lanzar la aplicación Gradio if __name__ == "__main__": pipeline = TrellisImageTo3DPipeline.from_pretrained("cavargas10/TRELLIS") pipeline.cuda() try: pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))) # Precargar rembg except: pass demo.launch(show_error=True)