hg_app.py

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--port', type=int, default=8080)
parser.add_argument('--cache-path', type=str, default='gradio_cache')
parser.add_argument('--enable_t23d', default=True)
parser.add_argument('--local', action="store_true")
args = parser.parse_args()
print(f"Running on {'local' if args.local else 'huggingface'}")

if not args.local:
    import os
    import spaces
    import subprocess
    import sys
    import shlex
    os.system("cd /home/user/app/hy3dgen/texgen/differentiable_renderer/ && bash compile_mesh_painter.sh")
    subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)    
    IP = "0.0.0.0"
    PORT = 7860
else:
    IP = "0.0.0.0"
    PORT = 8080
    class spaces:
        class GPU:
            def __init__(self, duration=60):
                self.duration = duration
            def __call__(self, func):
                return func 

import os
import shutil
import time
from glob import glob
from pathlib import Path
from PIL import Image
from datetime import datetime
import uuid
import gradio as gr
import torch
import uvicorn
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles

# Importar LitModel3D para renderizado eficiente
from gradio_litmodel3d import LitModel3D

def start_session(req: gr.Request):
    save_folder = os.path.join(SAVE_DIR, str(req.session_hash))
    os.makedirs(save_folder, exist_ok=True)

def end_session(req: gr.Request):
    save_folder = os.path.join(SAVE_DIR, str(req.session_hash))
    shutil.rmtree(save_folder)

def get_example_txt_list():
    print('Loading example txt list ...')
    txt_list = list()
    for line in open('./assets/example_prompts.txt'):
        txt_list.append(line.strip())
    return txt_list

def export_mesh(mesh, save_folder):
    path = os.path.join(save_folder, 'textured_mesh.glb')
    mesh.export(path)
    return path  # Retorna directamente la ruta del GLB

@spaces.GPU(duration=150)
def generation_all(
    caption: str,
    steps: int,
    guidance_scale: float,
    seed: int,
    octree_resolution: int,
    req: gr.Request,
):
    mesh, save_folder, image = _gen_shape(
        caption=caption,
        image=None,
        steps=steps,
        guidance_scale=guidance_scale,
        seed=seed,
        octree_resolution=octree_resolution,
        check_box_rembg=True,
        req=req
    )
    textured_mesh = texgen_worker(mesh, image)
    glb_path = export_mesh(textured_mesh, save_folder)
    torch.cuda.empty_cache()
    return glb_path  # Solo retorna la ruta para el visor

@spaces.GPU(duration=100)
def _gen_shape(
    caption: str,
    image: Image.Image,
    steps: int,
    guidance_scale: float,
    seed: int,
    octree_resolution: int,
    check_box_rembg: bool,
    req: gr.Request,
):
    if caption: print('prompt is', caption)
    save_folder = os.path.join(SAVE_DIR, str(req.session_hash)) 
    os.makedirs(save_folder, exist_ok=True)
    time_meta = {}
    start_time_0 = time.time()

    # Text-to-Image
    start_time = time.time()
    image = t2i_worker(caption)
    time_meta['text2image'] = time.time() - start_time
    image.save(os.path.join(save_folder, 'input.png'))

    # Remove background
    image = rmbg_worker(image.convert('RGB'))
    image.save(os.path.join(save_folder, 'rembg.png'))

    # Image-to-3D
    generator = torch.Generator().manual_seed(int(seed))
    mesh = i23d_worker(
        image=image,
        num_inference_steps=steps,
        guidance_scale=guidance_scale,
        generator=generator,
        octree_resolution=octree_resolution
    )[0]
    mesh = FloaterRemover()(mesh)
    mesh = DegenerateFaceRemover()(mesh)
    mesh = FaceReducer()(mesh)
    torch.cuda.empty_cache()
    
    return mesh, save_folder, image

def build_app():
    with gr.Blocks(theme=gr.themes.Base(), title='Hunyuan-3D-2.0') as demo:
        gr.Markdown("""
        # UTPL - Conversión de Texto 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 [Hunyuan3D v 2.0](https://github.com/tencent/Hunyuan3D-2) (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(scale=2):
                with gr.Tab('Text Prompt'):
                    caption = gr.Textbox(label='Text Prompt',
                                         placeholder='Enter a text description for 3D generation',
                                         info='Example: A cute cartoon rabbit with fluffy fur')
                with gr.Accordion('Advanced Options', open=False):
                    num_steps = gr.Slider(20, 50, 50, 1, label='Inference Steps')
                    octree_resolution = gr.Dropdown([256, 384, 512], 256, label='Octree Resolution')
                    cfg_scale = gr.Number(5.5, label='Guidance Scale')
                    seed = gr.Slider(0, 1e7, 1234, label='Seed')
                with gr.Group():
                    btn_all = gr.Button(value='Generate', variant='primary')

            with gr.Column(scale=5):
                # Usar LitModel3D para renderizado directo
                model_viewer = LitModel3D(
                    label="3D Model Viewer",
                    exposure=10.0,
                    height=400
                )
        
        # Conectar el botón al visor LitModel3D
        btn_all.click(
            generation_all,
            inputs=[
                caption,
                num_steps,
                cfg_scale,
                seed,
                octree_resolution,
            ],
            outputs=[model_viewer]  # Solo salida al visor
        )

    return demo

if __name__ == '__main__':
    CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
    SAVE_DIR = os.path.join(CURRENT_DIR, args.cache_path)
    os.makedirs(SAVE_DIR, exist_ok=True)
    example_ts = get_example_txt_list()
    
    from hy3dgen.texgen import Hunyuan3DPaintPipeline
    texgen_worker = Hunyuan3DPaintPipeline.from_pretrained('cavargas10/Hunyuan3D-2v')
    
    if args.enable_t23d:
        from hy3dgen.text2image import HunyuanDiTPipeline
        t2i_worker = HunyuanDiTPipeline('Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled')
    
    from hy3dgen.shapegen import FaceReducer, FloaterRemover, DegenerateFaceRemover
    from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline
    from hy3dgen.rembg import BackgroundRemover
    rmbg_worker = BackgroundRemover()
    i23d_worker = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained('cavargas10/Hunyuan3D-2v')
    
    app = FastAPI()
    app.mount("/static", StaticFiles(directory=SAVE_DIR), name="static")
    demo = build_app()
    demo.queue(max_size=10)
    app = gr.mount_gradio_app(app, demo, path="/")
    uvicorn.run(app, host=IP, port=PORT)