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() def initialize_models(device): try: print("Initializing models...") # 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 # CUDA 메모리 관리 설정 torch.cuda.empty_cache() torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.benchmark = True # 환경 변수 설정 # 환경 변수 설정 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]: print(f"Starting image_to_3d with trial_id: {trial_id}") if not trial_id or trial_id.strip() == "": print("Error: No trial_id provided") return None, None try: # CUDA 메모리 초기화 if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() gc.collect() if randomize_seed: seed = np.random.randint(0, MAX_SEED) image_path = f"{TMP_DIR}/{trial_id}.png" print(f"Looking for image at: {image_path}") if not os.path.exists(image_path): print(f"Error: Image file not found at {image_path}") return None, None image = Image.open(image_path) print(f"Successfully loaded image with size: {image.size}") # 이미지 크기 제한 max_size = 512 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) print(f"Resized image to: {image.size}") # GPU 작업 시작 with torch.inference_mode(): try: # 모델을 GPU로 이동 g.trellis_pipeline.to('cuda') torch.cuda.synchronize() # 3D 생성 outputs = g.trellis_pipeline.run( image, seed=seed, formats=["gaussian", "mesh"], preprocess_image=False, sparse_structure_sampler_params={ "steps": min(ss_sampling_steps, 12), "cfg_strength": ss_guidance_strength, }, slat_sampler_params={ "steps": min(slat_sampling_steps, 12), "cfg_strength": slat_guidance_strength, }, ) torch.cuda.synchronize() # 비디오 렌더링 video = render_utils.render_video( outputs['gaussian'][0], num_frames=60, resolution=512 )['color'] torch.cuda.synchronize() video_geo = render_utils.render_video( outputs['mesh'][0], num_frames=60, resolution=512 )['normal'] torch.cuda.synchronize() # CPU로 데이터 이동 video = [v.cpu().numpy() if torch.is_tensor(v) else v for v in video] video_geo = [v.cpu().numpy() if torch.is_tensor(v) else v for v in video_geo] 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" os.makedirs(os.path.dirname(video_path), exist_ok=True) 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') if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() gc.collect() except Exception as e: print(f"Error in image_to_3d: {str(e)}") if hasattr(g.trellis_pipeline, 'to'): g.trellis_pipeline.to('cpu') if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() gc.collect() return None, None def clear_gpu_memory(): """GPU 메모리를 정리하는 유틸리티 함수""" try: if torch.cuda.is_available(): with torch.cuda.device('cuda'): torch.cuda.empty_cache() torch.cuda.synchronize() gc.collect() except Exception as e: print(f"Error clearing 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]: try: # GPU 메모리 정리 clear_gpu_memory() # 상태 언패킹 gs, mesh, trial_id = unpack_state(state) # GLB 변환 전 유효성 검사 if gs is None or mesh is None: print("Error: Invalid gaussian or mesh data") return None, None # GLB 변환 with torch.inference_mode(): try: # 모든 텐서를 CUDA로 이동 (gradient 불필요) device = torch.device('cuda:0') # Gaussian 텐서들을 변환 for attr_name in ['_xyz', '_features_dc', '_scaling', '_rotation', '_opacity']: if hasattr(gs, attr_name): tensor = getattr(gs, attr_name) if torch.is_tensor(tensor): new_tensor = tensor.detach().clone().float().to(device) setattr(gs, attr_name, new_tensor) # Mesh 텐서들을 변환 if hasattr(mesh, 'vertices') and torch.is_tensor(mesh.vertices): mesh.vertices = mesh.vertices.detach().clone().float().to(device) if hasattr(mesh, 'faces') and torch.is_tensor(mesh.faces): mesh.faces = mesh.faces.detach().clone().long().to(device) # 추가 속성 확인 및 변환 (gradient 불필요) for attr_name in dir(mesh): if attr_name.startswith('_'): continue attr = getattr(mesh, attr_name) if torch.is_tensor(attr): if attr.dtype in [torch.float32, torch.float64]: setattr(mesh, attr_name, attr.to(device)) else: setattr(mesh, attr_name, attr.to(device)) print("Device and gradient check before GLB conversion:") print(f"Gaussian xyz device: {gs._xyz.device}, requires_grad: {gs._xyz.requires_grad}") print(f"Mesh vertices device: {mesh.vertices.device}, requires_grad: {mesh.vertices.requires_grad}") # GLB 변환 glb = postprocessing_utils.to_glb( gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=True ) except Exception as e: print(f"Error during GLB conversion: {str(e)}") # 디바이스와 gradient 정보 출력 if hasattr(gs, '_xyz'): print(f"Gaussian xyz device: {gs._xyz.device}, requires_grad: {gs._xyz.requires_grad}") if hasattr(mesh, 'vertices'): print(f"Mesh vertices device: {mesh.vertices.device}, requires_grad: {mesh.vertices.requires_grad}") return None, None if glb is None: print("Error: GLB conversion failed") return None, None # 파일 저장 glb_path = f"{TMP_DIR}/{trial_id}.glb" try: glb.export(glb_path) if not os.path.exists(glb_path): print(f"Error: GLB file was not created at {glb_path}") return None, None except Exception as e: print(f"Error saving GLB file: {str(e)}") return None, None print(f"Successfully created GLB file at: {glb_path}") return glb_path, glb_path except Exception as e: print(f"Error in extract_glb: {str(e)}") return None, None finally: # 정리 작업 clear_gpu_memory() 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 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() with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("""## Craft3D""") # 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(): 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