app.py

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)