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, ImageOps
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import render_utils, postprocessing_utils
import torch
import torchvision.transforms.functional as TF
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - HF_SPACE_BOCETO - %(levelname)s - %(message)s')

style_list = [
    {
        "name": "(No style)",
        "prompt": "{prompt}",
        "negative_prompt": "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
    },
    {
        "name": "Cinematic",
        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
    },
    {
        "name": "3D Model",
        "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
    },
    {
        "name": "Anime",
        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
    },
]
styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "3D Model"
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)

def start_session(req: gr.Request):
    session_hash = str(req.session_hash)
    user_dir = os.path.join(TMP_DIR, session_hash)
    logging.info(f"START SESSION: Creando directorio para la sesión {session_hash} en {user_dir}")
    os.makedirs(user_dir, exist_ok=True)

def end_session(req: gr.Request):
    session_hash = str(req.session_hash)
    user_dir = os.path.join(TMP_DIR, session_hash)
    logging.info(f"END SESSION: Intentando eliminar el directorio de la sesión {session_hash} en {user_dir}")
    if os.path.exists(user_dir):
        try:
            shutil.rmtree(user_dir)
            logging.info(f"Directorio de la sesión {session_hash} eliminado correctamente.")
        except Exception as e:
            logging.error(f"Error al eliminar el directorio de la sesión {session_hash}: {e}")
    else:
        logging.warning(f"El directorio de la sesión {session_hash} no fue encontrado al intentar eliminarlo. Es posible que ya haya sido limpiado.")

def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]:
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    return p.replace("{prompt}", positive), n + negative

def get_seed(randomize_seed: bool, seed: int) -> int:
    new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed
    logging.info(f"Usando seed: {new_seed}")
    return new_seed

@spaces.GPU
def preprocess_image(
    image: Image.Image,
    prompt: str = "",
    negative_prompt: str = "",
    style_name: str = "",
    num_steps: int = 25,
    guidance_scale: float = 5,
    controlnet_conditioning_scale: float = 1.0,
    req: gr.Request = None,
) -> str:
    session_hash = str(req.session_hash)
    user_dir = os.path.join(TMP_DIR, session_hash)
    logging.info(f"[{session_hash}] Iniciando preprocess_image con prompt: '{prompt[:50]}...'")

    if image is None:
        logging.error(f"[{session_hash}] La entrada de imagen es nula.")
        raise ValueError("La imagen de entrada no puede estar vacía.")

    input_image = image
    
    width, height = input_image.size
    ratio = np.sqrt(1024.0 * 1024.0 / (width * height))
    new_width, new_height = int(width * ratio), int(height * ratio)
    input_image = input_image.resize((new_width, new_height))
    
    if input_image.mode == 'RGBA':
        r, g, b, a = input_image.split()
        rgb_image = Image.merge('RGB', (r, g, b))
        inverted_image = ImageOps.invert(rgb_image)
        inverted_image.putalpha(a)
        input_image = inverted_image
    else:
        input_image = ImageOps.invert(input_image.convert('RGB'))

    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
    
    output_image = pipe_control(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=input_image,
        num_inference_steps=num_steps,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guidance_scale=guidance_scale,
        width=new_width,
        height=new_height,
    ).images[0]
    
    processed_image_path = os.path.join(user_dir, 'processed_image.png')
    output_image.save(processed_image_path)
    logging.info(f"[{session_hash}] Imagen preprocesada y guardada en: {processed_image_path}")
    
    return processed_image_path

@spaces.GPU
def image_to_3d(
    image_path: str,
    seed: int,
    ss_guidance_strength: float,
    ss_sampling_steps: int,
    slat_guidance_strength: float,
    slat_sampling_steps: int,
    req: gr.Request,
) -> Tuple[dict, str]:
    session_hash = str(req.session_hash)
    user_dir = os.path.join(TMP_DIR, session_hash)
    logging.info(f"[{session_hash}] Iniciando image_to_3d desde la imagen: {image_path}")
    processed_image = pipeline.preprocess_image(Image.open(image_path))
    outputs = pipeline.run(
        processed_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},
    )
    logging.info(f"[{session_hash}] Generación del modelo completada. Renderizando video...")
    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()
    logging.info(f"[{session_hash}] Video renderizado y estado empaquetado. Devolviendo: {video_path}")
    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]:
    session_hash = str(req.session_hash)
    user_dir = os.path.join(TMP_DIR, session_hash)
    logging.info(f"[{session_hash}] Iniciando extract_glb...")
    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()
    logging.info(f"[{session_hash}] GLB extraído. Devolviendo: {glb_path}")
    return glb_path, glb_path

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, str]:
    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

@spaces.GPU
def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
    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

with gr.Blocks(delete_cache=(600, 600)) as demo:
    gr.Markdown("""
    # UTPL - Conversión de Boceto 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.
    ---
    **Modelos Utilizados:**
    - **ControlNet Scribble:** `xinsir/controlnet-scribble-sdxl-1.0`
    - **Stable Diffusion Base:** `sd-community/sdxl-flash`
    - **VAE:** `madebyollin/sdxl-vae-fp16-fix`
    """)
    with gr.Row():
        with gr.Column():
            with gr.Column():
                # --- ¡MODIFICADO! Cambiamos ImageEditor por Image ---
                image_prompt = gr.Image(label="Input sketch", type="pil", image_mode="RGBA", height=512)
                
                with gr.Row():
                    sketch_btn = gr.Button("Process Sketch")
                    generate_btn = gr.Button("Generate 3D")
                with gr.Row():
                    prompt = gr.Textbox(label="Prompt")
                    style = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
                with gr.Accordion(label="Generation Settings", open=False):
                    with gr.Tab(label="Sketch-to-Image Generation"):
                        negative_prompt = gr.Textbox(label="Negative prompt")
                        num_steps = gr.Slider(1, 20, label="Number of steps", value=8, step=1)
                        guidance_scale = gr.Slider(0.1, 10.0, label="Guidance scale", value=5, step=0.1)
                        controlnet_conditioning_scale = gr.Slider(0.5, 5.0, label="ControlNet Conditioning Scale", value=0.85, step=0.01)
                    with gr.Tab(label="3D Generation"):
                        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)
            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)
        with gr.Column():
            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
            image_prompt_processed = gr.Image(label="Processed Sketch", interactive=False, type="filepath", height=512)
            model_output = gr.Model3D(label="Extracted GLB/Gaussian", height=300)
            with gr.Row():
                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)  

    output_buf = gr.State()

    demo.load(start_session)
    demo.unload(end_session)

    sketch_btn.click(
        preprocess_image,
        inputs=[image_prompt, prompt, negative_prompt, style, num_steps, guidance_scale, controlnet_conditioning_scale],
        outputs=[image_prompt_processed],
        api_name="preprocess_image" 
    )

    generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
    ).then(
        image_to_3d,
        inputs=[image_prompt_processed, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
        outputs=[output_buf, video_output],
        api_name="image_to_3d"
    )

    generate_btn.click(
        lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
        outputs=[extract_glb_btn, extract_gs_btn],
    )

    video_output.clear(
        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]),
        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],
        api_name="extract_glb" 
    )
    
    extract_glb_btn.click(
        lambda: gr.Button(interactive=True),
        outputs=[download_glb]
    )
    
    extract_gs_btn.click(
        extract_gaussian,
        inputs=[output_buf],
        outputs=[model_output, download_gs],
        api_name="extract_gaussian" 
    )

    extract_gs_btn.click(
        lambda: gr.Button(interactive=True),
        outputs=[download_gs]
    )

    model_output.clear(
        lambda: gr.Button(interactive=False),
        outputs=[download_glb],
    )

if __name__ == "__main__":
    pipeline = TrellisImageTo3DPipeline.from_pretrained("cavargas10/TRELLIS")
    pipeline.cuda()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    controlnet = ControlNetModel.from_pretrained("xinsir/controlnet-scribble-sdxl-1.0", torch_dtype=torch.float16)
    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
    pipe_control = StableDiffusionXLControlNetPipeline.from_pretrained("sd-community/sdxl-flash", controlnet=controlnet, vae=vae, torch_dtype=torch.float16)
    pipe_control.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_control.scheduler.config)
    pipe_control.to(device)
    try:
        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))
    except:
        pass
    demo.launch(show_error=True)