Update app.py

app.py

@@ -3,19 +3,46 @@ import spaces
 from gradio_litmodel3d import LitModel3D
 import os
 import shutil
-import torch
 import numpy as np
+import torch
 import imageio
 from easydict import EasyDict as edict
-from PIL import Image
+from PIL import Image, ImageOps
 from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
+from typing import List, Tuple, Literal
+from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
+from diffusers import EulerAncestralDiscreteScheduler
+from controlnet_aux import PidiNetDetector, HEDdetector
+
+os.environ['SPCONV_ALGO'] = 'native'
+
+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",
+    },
+    # ... (otros estilos de la lista original)
+]
 
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "(No style)"
 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 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 start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
@@ -24,8 +51,35 @@ def end_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     shutil.rmtree(user_dir, ignore_errors=True)
 
-def preprocess_image(image: Image.Image) -> Image.Image:
-    return pipeline.preprocess_image(image)
+@spaces.GPU
+def preprocess_image(image: Image.Image,
+                    prompt: str = "",
+                    negative_prompt: str = "",
+                    style_name: str = DEFAULT_STYLE_NAME,
+                    num_steps: int = 25,
+                    guidance_scale: float = 5,
+                    controlnet_conditioning_scale: float = 1.0) -> Image.Image:
+    
+    width, height = image.size
+    ratio = np.sqrt(1024 * 1024 / (width * height))
+    new_size = (int(width * ratio), int(height * ratio))
+    image = image.resize(new_size)
+    image = ImageOps.invert(image.convert("L")).convert("RGB")
+    
+    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
+    
+    output = pipe_control(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=image,
+        num_inference_steps=num_steps,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        guidance_scale=guidance_scale,
+        width=new_size[0],
+        height=new_size[1]
+    ).images[0]
+    
+    return pipeline.preprocess_image(output)
 
 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
     return {
@@ -45,12 +99,6 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
 
 def unpack_state(state: dict) -> Tuple[Gaussian, edict]:
     gs = Gaussian(**state['gaussian'])
-    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'),
@@ -61,65 +109,169 @@ 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(image: Image.Image, seed: int, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int, req: gr.Request) -> Tuple[dict, str]:
+def image_to_3d(
+    image: Image.Image,
+    seed: int,
+    ss_guidance_strength: float,
+    ss_sampling_steps: int,
+    slat_guidance_strength: float,
+    slat_sampling_steps: int,
+    req: gr.Request,
+) -> Tuple[dict, str]:
+    
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     outputs = pipeline.run(
         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},
+        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,
+        },
     )
 
     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 = [np.concatenate([v, g], axis=1) for v, g in zip(video, video_geo)]
+    
     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
+    
+    return pack_state(outputs['gaussian'][0], outputs['mesh'][0]), video_path
 
 @spaces.GPU(duration=90)
-def extract_glb(state: dict, mesh_simplify: float, texture_size: int, req: gr.Request) -> Tuple[str, str]:
+def extract_glb(
+    state: dict,
+    mesh_simplify: float,
+    texture_size: int,
+    req: gr.Request,
+) -> 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 = 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
+    return glb_path
 
 with gr.Blocks() as demo:
-    gr.Markdown("""
-    # Conversión de Imágen a 3D
-    """)
-    
-    image_prompt = gr.Image(label="Input Image", type="pil")
-    seed = gr.Slider(0, MAX_SEED, label="Seed", value=0)
-    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-    generate_btn = gr.Button("Generate 3D Asset")
-    video_output = gr.Video(label="3D Preview")
-    model_output = LitModel3D(label="3D Model Viewer")
-    extract_glb_btn = gr.Button("Export GLB")
-    download_glb = gr.DownloadButton(label="Download GLB")
-    output_buf = gr.State()
+    gr.Markdown("# Sketch-to-3D con TRELLIS")
     
+    with gr.Row():
+        with gr.Column():
+            image_prompt = gr.Image(
+                label="Boceto",
+                type="pil",
+                image_mode="RGBA",
+                height=512,
+                tool="sketch"
+            )
+            
+            with gr.Accordion("Ajustes de Generación", open=False):
+                prompt = gr.Textbox(label="Prompt")
+                style = gr.Dropdown(label="Estilo", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
+                negative_prompt = gr.Textbox(label="Negative Prompt")
+                
+                with gr.Group():
+                    gr.Markdown("#### Etapa 1: Estructura")
+                    ss_guidance_strength = gr.Slider(0.0, 10.0, 7.5, step=0.1, label="Guidance")
+                    ss_sampling_steps = gr.Slider(1, 50, 12, step=1, label="Pasos")
+                
+                with gr.Group():
+                    gr.Markdown("#### Etapa 2: Detalle")
+                    slat_guidance_strength = gr.Slider(0.0, 10.0, 3.0, step=0.1, label="Guidance")
+                    slat_sampling_steps = gr.Slider(1, 50, 12, step=1, label="Pasos")
+            
+            generate_btn = gr.Button("Generar 3D", variant="primary")
+            
+            with gr.Accordion("Exportar GLB", open=False):
+                mesh_simplify = gr.Slider(0.9, 0.98, 0.95, step=0.01, label="Simplificar")
+                texture_size = gr.Slider(512, 2048, 1024, step=512, label="Tamaño Textura")
+                export_glb_btn = gr.Button("Exportar GLB", interactive=False)
+        
+        with gr.Column():
+            video_output = gr.Video(label="Vista 3D", autoplay=True, loop=True, height=300)
+            model_viewer = LitModel3D(label="Visor 3D", height=400)
+            download_glb = gr.DownloadButton("Descargar GLB", interactive=False)
+
+    output_state = gr.State()
+
     generate_btn.click(
-        get_seed, inputs=[randomize_seed, seed], outputs=[seed]
+        get_seed,
+        inputs=[gr.Checkbox(value=True, label="Semilla Aleatoria"), gr.Number(0, visible=False)],
+        outputs=[gr.Number(0, visible=False)]
+    ).then(
+        preprocess_image,
+        inputs=[
+            image_prompt,
+            prompt,
+            negative_prompt,
+            style,
+            gr.Slider(1, 20, 8, step=1, label="Pasos SDXL"),
+            gr.Slider(0.1, 10.0, 5.0, step=0.1, label="Guidance SDXL"),
+            gr.Slider(0.5, 5.0, 0.85, step=0.01, label="ControlNet Strength")
+        ],
+        outputs=[gr.Image(label="Imagen Procesada")]
     ).then(
-        image_to_3d, inputs=[image_prompt, seed, 7.5, 12, 3.0, 12], outputs=[output_buf, video_output]
+        image_to_3d,
+        inputs=[
+            gr.Image(visible=False),
+            gr.Number(0),
+            ss_guidance_strength,
+            ss_sampling_steps,
+            slat_guidance_strength,
+            slat_sampling_steps
+        ],
+        outputs=[output_state, video_output]
     ).then(
-        lambda: gr.Button(interactive=True), outputs=[extract_glb_btn]
+        lambda: gr.update(interactive=True),
+        outputs=[export_glb_btn]
     )
-    
-    extract_glb_btn.click(
-        extract_glb, inputs=[output_buf, 0.95, 1024], outputs=[model_output, download_glb]
+
+    export_glb_btn.click(
+        extract_glb,
+        inputs=[output_state, mesh_simplify, texture_size],
+        outputs=[model_viewer, download_glb]
     ).then(
-        lambda: gr.Button(interactive=True), outputs=[download_glb]
+        lambda: gr.update(interactive=True),
+        outputs=[download_glb]
     )
 
 if __name__ == "__main__":
-    pipeline = TrellisImageTo3DPipeline.from_pretrained("cavargas10/TRELLIS").cuda()
-    demo.launch()
+    # Inicializar pipelines
+    pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large").cuda()
+    
+    # ControlNet SDXL
+    controlnet = ControlNetModel.from_pretrained(
+        "xinsir/controlnet-scribble-sdxl-1.0",
+        torch_dtype=torch.float16
+    ).cuda()
+    
+    vae = AutoencoderKL.from_pretrained(
+        "madebyollin/sdxl-vae-fp16-fix",
+        torch_dtype=torch.float16
+    ).cuda()
+    
+    pipe_control = StableDiffusionXLControlNetPipeline.from_pretrained(
+        "sd-community/sdxl-flash",
+        controlnet=controlnet,
+        vae=vae,
+        torch_dtype=torch.float16
+    ).cuda()
+    
+    pipe_control.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_control.scheduler.config)
+    
+    demo.launch()