Update app.py

app.py

@@ -1,622 +0,0 @@
-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