app.py

import gradio as gr
import spaces
from gradio_litmodel3d import LitModel3D
import os
import time
from os import path
import shutil
from datetime import datetime
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
import torch
import numpy as np
import imageio
import uuid
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
from diffusers import FluxPipeline
from typing import Tuple, Dict, Any  # Tuple import 추가
# 파일 상단의 import 문
import transformers
from transformers import pipeline as transformers_pipeline
from transformers import Pipeline
import gc  # 파일 상단에 추가

# 전역 변수 초기화
class GlobalVars:
    def __init__(self):
        self.translator = None
        self.trellis_pipeline = None
        self.flux_pipe = None

g = GlobalVars()

# 파일 상단에 추가
torch.backends.cudnn.benchmark = False  # 메모리 사용량 감소
torch.backends.cudnn.deterministic = True
torch.cuda.set_per_process_memory_fraction(0.7)  # GPU 메모리 사용량 제한
        
def initialize_models(device):
    try:
        print("Initializing models...")
        g.translator = transformers_pipeline(
            "translation", 
            model="Helsinki-NLP/opus-mt-ko-en",
            device=device
        )
        print("Model initialization completed successfully")
        
        # 3D 생성 파이프라인
        g.trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
            "JeffreyXiang/TRELLIS-image-large"
        )
        print("TrellisImageTo3DPipeline loaded successfully")
        
        # 이미지 생성 파이프라인
        print("Loading flux_pipe...")
        g.flux_pipe = FluxPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            torch_dtype=torch.bfloat16,
            device_map="balanced"
        )
        print("FluxPipeline loaded successfully")
        
        # Hyper-SD LoRA 로드
        print("Loading LoRA weights...")
        lora_path = hf_hub_download(
            "ByteDance/Hyper-SD",
            "Hyper-FLUX.1-dev-8steps-lora.safetensors",
            use_auth_token=HF_TOKEN
        )
        g.flux_pipe.load_lora_weights(lora_path)
        g.flux_pipe.fuse_lora(lora_scale=0.125)
        print("LoRA weights loaded successfully")
        
        # 번역기 초기화
        print("Initializing translator...")
        g.translator = transformers_pipeline(
            "translation", 
            model="Helsinki-NLP/opus-mt-ko-en",
            device=device
        )
        print("Model initialization completed successfully")
        
    except Exception as e:
        print(f"Error during model initialization: {str(e)}")
        raise
    

# 환경 변수 설정
# 환경 변수 설정
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
os.environ['SPCONV_ALGO'] = 'native'
os.environ['SPARSE_BACKEND'] = 'native'
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
os.environ['XFORMERS_FORCE_DISABLE_TRITON'] = '1'
os.environ['XFORMERS_ENABLE_FLASH_ATTENTION'] = '1'
os.environ['TORCH_CUDA_MEMORY_ALLOCATOR'] = 'native'
os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'

# CUDA 초기화 방지
torch.set_grad_enabled(False)

# Hugging Face 토큰 설정
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN is None:
    raise ValueError("HF_TOKEN environment variable is not set")

MAX_SEED = np.iinfo(np.int32).max
TMP_DIR = "/tmp/Trellis-demo"
os.makedirs(TMP_DIR, exist_ok=True)

# Setup and initialization code
cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
PERSISTENT_DIR = os.environ.get("PERSISTENT_DIR", ".")
gallery_path = path.join(PERSISTENT_DIR, "gallery")

os.environ["TRANSFORMERS_CACHE"] = cache_path
os.environ["HF_HUB_CACHE"] = cache_path
os.environ["HF_HOME"] = cache_path
os.environ['SPCONV_ALGO'] = 'native'

torch.backends.cuda.matmul.allow_tf32 = True



class timer:
    def __init__(self, method_name="timed process"):
        self.method = method_name
    def __enter__(self):
        self.start = time.time()
        print(f"{self.method} starts")
    def __exit__(self, exc_type, exc_val, exc_tb):
        end = time.time()
        print(f"{self.method} took {str(round(end - self.start, 2))}s")

def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
    if image is None:
        print("Error: Input image is None")
        return "", None
        
    try:
        if g.trellis_pipeline is None:
            print("Error: trellis_pipeline is not initialized")
            return "", None
            
        # webp 이미지를 RGB로 변환
        if isinstance(image, str) and image.endswith('.webp'):
            image = Image.open(image).convert('RGB')
        elif isinstance(image, Image.Image):
            image = image.convert('RGB')
            
        trial_id = str(uuid.uuid4())
        processed_image = g.trellis_pipeline.preprocess_image(image)
        if processed_image is not None:
            save_path = f"{TMP_DIR}/{trial_id}.png"
            processed_image.save(save_path)
            print(f"Saved processed image to: {save_path}")
            return trial_id, processed_image
        else:
            print("Error: Processed image is None")
            return "", None
    except Exception as e:
        print(f"Error in image preprocessing: {str(e)}")
        return "", None
        
def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> 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(),
        },
        'trial_id': trial_id,
    }
    
    
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, state['trial_id']

@spaces.GPU
def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_strength: float, 
                ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int) -> Tuple[dict, str]:
    try:
        # 초기 메모리 정리
        clear_gpu_memory()
        
        if not trial_id or trial_id.strip() == "":
            return None, None
            
        image_path = f"{TMP_DIR}/{trial_id}.png"
        if not os.path.exists(image_path):
            return None, None
            
        image = Image.open(image_path)
        
        # 이미지 크기 제한 강화
        max_size = 384  # 더 작은 크기로 제한
        if max(image.size) > max_size:
            ratio = max_size / max(image.size)
            new_size = tuple(int(dim * ratio) for dim in image.size)
            image = image.resize(new_size, Image.LANCZOS)
        
        with torch.inference_mode():
            try:
                # 파이프라인을 GPU로 이동
                g.trellis_pipeline.to('cuda')
                
                outputs = g.trellis_pipeline.run(
                    image,
                    seed=seed,
                    formats=["gaussian", "mesh"],
                    preprocess_image=False,
                    sparse_structure_sampler_params={
                        "steps": min(ss_sampling_steps, 8),
                        "cfg_strength": ss_guidance_strength
                    },
                    slat_sampler_params={
                        "steps": min(slat_sampling_steps, 8),
                        "cfg_strength": slat_guidance_strength
                    }
                )
                
                # 중간 메모리 정리
                clear_gpu_memory()
                
                # 비디오 렌더링 최적화
                video = render_utils.render_video(
                    outputs['gaussian'][0],
                    num_frames=30,
                    resolution=384
                )['color']
                
                video_geo = render_utils.render_video(
                    outputs['mesh'][0],
                    num_frames=30,
                    resolution=384
                )['normal']
                
                # tensor를 numpy로 변환
                if torch.is_tensor(video[0]):
                    video = [v.cpu().numpy() if torch.is_tensor(v) else v for v in video]
                if torch.is_tensor(video_geo[0]):
                    video_geo = [v.cpu().numpy() if torch.is_tensor(v) else v for v in video_geo]
                
                clear_gpu_memory()
                
                # 비디오 생성 및 저장
                video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
                new_trial_id = str(uuid.uuid4())
                video_path = f"{TMP_DIR}/{new_trial_id}.mp4"
                imageio.mimsave(video_path, video, fps=15)
                
                state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], new_trial_id)
                return state, video_path
                
            finally:
                # 정리 작업
                g.trellis_pipeline.to('cpu')
                clear_gpu_memory()
                
    except Exception as e:
        print(f"Error in image_to_3d: {str(e)}")
        if hasattr(g.trellis_pipeline, 'to'):
            g.trellis_pipeline.to('cpu')
        clear_gpu_memory()
        return None, None

def clear_gpu_memory():
    """GPU 메모리를 더 철저하게 정리하는 함수"""
    try:
        if torch.cuda.is_available():
            # 모든 GPU 캐시 정리
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
            
            # 사용하지 않는 캐시된 메모리 해제
            for i in range(torch.cuda.device_count()):
                with torch.cuda.device(i):
                    torch.cuda.empty_cache()
                    torch.cuda.ipc_collect()
                    
        # Python 가비지 컬렉터 실행
        gc.collect()
    except Exception as e:
        print(f"Error in clear_gpu_memory: {e}")

def move_to_device(model, device):
    """모델을 안전하게 디바이스로 이동하는 함수"""
    try:
        if hasattr(model, 'to'):
            clear_gpu_memory()
            model.to(device)
            if device == 'cuda':
                torch.cuda.synchronize()
            clear_gpu_memory()
    except Exception as e:
        print(f"Error moving model to {device}: {str(e)}")
        
@spaces.GPU
def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[str, str]:
    """
    3D 모델에서 GLB 파일 추출
    """
    try:
        gs, mesh, trial_id = unpack_state(state)
        glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
        glb_path = f"{TMP_DIR}/{trial_id}.glb"
        glb.export(glb_path)
        return glb_path, glb_path
    except Exception as e:
        print(f"GLB 추출 중 오류 발생: {e}")
        return None, None



def activate_button() -> gr.Button:
    return gr.Button(interactive=True)


def deactivate_button() -> gr.Button:
    return gr.Button(interactive=False)

@spaces.GPU
def text_to_image(prompt: str, height: int, width: int, steps: int, scales: float, seed: int) -> Image.Image:
    try:
        # CUDA 메모리 초기화
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
        gc.collect()
        
        # 한글 감지 및 번역
        def contains_korean(text):
            return any(ord('가') <= ord(c) <= ord('힣') for c in text)
        
        if contains_korean(prompt):
            translated = g.translator(prompt)[0]['translation_text']
            prompt = translated
            print(f"Translated prompt: {prompt}")
        
        formatted_prompt = f"wbgmsst, 3D, {prompt}, white background"
        
        # 크기 제한
        height = min(height, 512)
        width = min(width, 512)
        steps = min(steps, 12)
        
        with torch.inference_mode():
            generated_image = g.flux_pipe(
                prompt=[formatted_prompt],
                generator=torch.Generator('cuda').manual_seed(int(seed)),
                num_inference_steps=int(steps),
                guidance_scale=float(scales),
                height=int(height),
                width=int(width),
                max_sequence_length=256
            ).images[0]
            
            if generated_image is not None:
                trial_id = str(uuid.uuid4())
                save_path = f"{TMP_DIR}/{trial_id}.png"
                generated_image.save(save_path)
                print(f"Saved generated image to: {save_path}")
                return generated_image
            else:
                print("Error: Generated image is None")
                return None
                
    except Exception as e:
        print(f"Error in image generation: {str(e)}")
        return None
    finally:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
        gc.collect()
    
css = """
footer {
    visibility: hidden;
}
"""

def periodic_cleanup():
    """주기적으로 실행될 메모리 정리 함수"""
    clear_gpu_memory()
    return None

with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.Markdown("""## Roblox3D GEN""")
    
    # Examples 이미지 로드
    example_dir = "assets/example_image/"
    example_images = []
    if os.path.exists(example_dir):
        for file in os.listdir(example_dir):
            if file.endswith('.webp'):
                example_images.append(os.path.join(example_dir, file))
    
    with gr.Row():
        with gr.Column():
            text_prompt = gr.Textbox(
                label="Text Prompt",
                placeholder="Describe what you want to create...",
                lines=3
            )
            
            # 이미지 프롬프트
            image_prompt = gr.Image(label="Image Prompt", image_mode="RGBA", type="pil", height=300)
            
            with gr.Accordion("Image Generation Settings", open=False):
                with gr.Row():
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=1152,
                        step=64,
                        value=1024
                    )
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=1152,
                        step=64,
                        value=1024
                    )
                
                with gr.Row():
                    steps = gr.Slider(
                        label="Inference Steps",
                        minimum=6,
                        maximum=25,
                        step=1,
                        value=8
                    )
                    scales = gr.Slider(
                        label="Guidance Scale",
                        minimum=0.0,
                        maximum=5.0,
                        step=0.1,
                        value=3.5
                    )
                
                seed = gr.Number(
                    label="Seed",
                    value=lambda: torch.randint(0, MAX_SEED, (1,)).item(),
                    precision=0
                )
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            
            generate_image_btn = gr.Button("Generate Image")
            
            with gr.Accordion("3D Generation Settings", open=False):
                ss_guidance_strength = gr.Slider(0.0, 10.0, label="Structure Guidance Strength", value=7.5, step=0.1)
                ss_sampling_steps = gr.Slider(1, 50, label="Structure Sampling Steps", value=12, step=1)
                slat_guidance_strength = gr.Slider(0.0, 10.0, label="Latent Guidance Strength", value=3.0, step=0.1)
                slat_sampling_steps = gr.Slider(1, 50, label="Latent Sampling Steps", value=12, step=1)

            generate_3d_btn = gr.Button("Generate 3D")
            
            with gr.Accordion("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():
            # 상단에 3d.mp4 자동재생 비디오 추가
            gr.Video(
                "3d.mp4",
                label="3D Asset Preview",
                autoplay=True,
                loop=True,
                height=300,
                width="100%"
            )
            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
            model_output = LitModel3D(label="Extracted GLB", exposure=20.0, height=300)
            download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
    
    trial_id = gr.Textbox(visible=False)
    output_buf = gr.State()

    # Examples 갤러리를 맨 아래로 이동
    if example_images:
        gr.Markdown("""### Example Images""")
        with gr.Row():
            gallery = gr.Gallery(
                value=example_images,
                label="Click an image to use it",
                show_label=True,
                elem_id="gallery",
                columns=11,  # 한 줄에 12개
                rows=3,      # 2줄
                height=400,  # 높이 조정
                allow_preview=True,
                object_fit="contain"  # 이미지 비율 유지
            )
            
            def load_example(evt: gr.SelectData):
                selected_image = Image.open(example_images[evt.index])
                trial_id_val, processed_image = preprocess_image(selected_image)
                return selected_image, trial_id_val
            
            gallery.select(
                load_example,
                None,
                [image_prompt, trial_id],
                show_progress=True
            )

    # Handlers
    generate_image_btn.click(
        text_to_image,
        inputs=[text_prompt, height, width, steps, scales, seed],
        outputs=[image_prompt]
    ).then(
        preprocess_image,
        inputs=[image_prompt],
        outputs=[trial_id, image_prompt]
    )

    # 나머지 핸들러들
    image_prompt.upload(
        preprocess_image,
        inputs=[image_prompt],
        outputs=[trial_id, image_prompt],
    )
    
    image_prompt.clear(
        lambda: '',
        outputs=[trial_id],
    )

    generate_3d_btn.click(
        image_to_3d,
        inputs=[trial_id, seed, randomize_seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
        outputs=[output_buf, video_output],
    ).then(
        activate_button,
        outputs=[extract_glb_btn],
    )

    video_output.clear(
        deactivate_button,
        outputs=[extract_glb_btn],
    )

    extract_glb_btn.click(
        extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size],
        outputs=[model_output, download_glb],
    ).then(
        activate_button,
        outputs=[download_glb],
    )

    model_output.clear(
        deactivate_button,
        outputs=[download_glb],
    )

if __name__ == "__main__":
    try:
        # CPU로 초기화
        device = "cpu"
        print(f"Using device: {device}")
        
        # 모델 초기화
        initialize_models(device)
        
        # 초기 이미지 전처리 테스트
        try:
            test_image = Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))
            if g.trellis_pipeline is not None:
                g.trellis_pipeline.preprocess_image(test_image)
            else:
                print("Warning: trellis_pipeline is None")
        except Exception as e:
            print(f"Warning: Initial preprocessing test failed: {e}")
            
        # Gradio 인터페이스 실행
        demo.queue()  # 큐 기능 활성화
        demo.launch(
            allowed_paths=[PERSISTENT_DIR, TMP_DIR],
            server_name="0.0.0.0",
            server_port=7860,
            show_error=True
        )
        
    except Exception as e:
        print(f"Error during initialization: {e}")
        raise