import gradio as gr import spaces from gradio_litmodel3d import LitModel3D import os 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 transformers import pipeline as translation_pipeline from diffusers import FluxPipeline from typing import * MAX_SEED = np.iinfo(np.int32).max TMP_DIR = "/tmp/Trellis-demo" os.makedirs(TMP_DIR, exist_ok=True) # GPU 메모리 관련 환경 변수 수정 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512' # A100에 맞게 증가 os.environ['CUDA_VISIBLE_DEVICES'] = '0' # 단일 GPU 사용 os.environ['CUDA_LAUNCH_BLOCKING'] = '0' # A100에서는 비동기 실행 허용 def initialize_models(): global pipeline, translator, flux_pipe try: import torch # L40S GPU 최적화 설정 torch.backends.cudnn.benchmark = True torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True print("Initializing Trellis pipeline...") pipeline = TrellisImageTo3DPipeline.from_pretrained( "JeffreyXiang/TRELLIS-image-large" ) if torch.cuda.is_available(): pipeline = pipeline.to("cuda") # 모델을 FP16으로 변환 for param in pipeline.parameters(): param.data = param.data.half() print("Initializing translator...") translator = translation_pipeline( "translation", model="Helsinki-NLP/opus-mt-ko-en", device="cuda" ) # Flux 파이프라인은 나중에 초기화 flux_pipe = None print("Models initialized successfully") return True except Exception as e: print(f"Model initialization error: {str(e)}") return False def get_flux_pipe(): """Flux 파이프라인을 필요할 때만 로드하는 함수""" global flux_pipe if flux_pipe is None: try: free_memory() flux_pipe = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", use_safetensors=True ).to("cuda") # FP16으로 변환 flux_pipe.to(torch.float16) except Exception as e: print(f"Error loading Flux pipeline: {e}") return None return flux_pipe def free_memory(): """강화된 메모리 정리 함수""" import gc import os # Python 가비지 컬렉션 gc.collect() # CUDA 메모리 정리 if torch.cuda.is_available(): torch.cuda.empty_cache() # 임시 파일 정리 tmp_dirs = ['/tmp/transformers_cache', '/tmp/torch_home', '/tmp/huggingface', '/tmp/cache'] for dir_path in tmp_dirs: if os.path.exists(dir_path): try: for file in os.listdir(dir_path): file_path = os.path.join(dir_path, file) if os.path.isfile(file_path): try: os.unlink(file_path) except: pass except: pass def setup_gpu_model(model): """GPU 설정이 필요한 모델을 처리하는 함수""" if torch.cuda.is_available(): model = model.to("cuda") return model def translate_if_korean(text): if any(ord('가') <= ord(char) <= ord('힣') for char in text): translated = translator(text)[0]['translation_text'] return translated return text def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]: try: if pipeline is None: raise Exception("Pipeline not initialized") trial_id = str(uuid.uuid4()) # 이미지가 너무 작은 경우 크기 조정 min_size = 64 if image.size[0] < min_size or image.size[1] < min_size: ratio = min_size / min(image.size) new_size = tuple(int(dim * ratio) for dim in image.size) image = image.resize(new_size, Image.LANCZOS) try: processed_image = pipeline.preprocess_image(image) if processed_image is None: raise Exception("Failed to process image") processed_image.save(f"{TMP_DIR}/{trial_id}.png") return trial_id, processed_image except Exception as e: print(f"Error in image preprocessing: {str(e)}") return None, None except Exception as e: print(f"Error in preprocess_image: {str(e)}") return None, 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'] 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): try: if randomize_seed: seed = np.random.randint(0, MAX_SEED) input_image = Image.open(f"{TMP_DIR}/{trial_id}.png") # L40S에 맞게 이미지 크기 제한 조정 max_size = 768 # L40S는 더 큰 이미지 처리 가능 if max(input_image.size) > max_size: ratio = max_size / max(input_image.size) input_image = input_image.resize( (int(input_image.size[0] * ratio), int(input_image.size[1] * ratio)), Image.LANCZOS ) if torch.cuda.is_available(): pipeline.to("cuda") with torch.cuda.amp.autocast(): # 자동 혼합 정밀도 사용 outputs = pipeline.run( input_image, seed=seed, formats=["gaussian", "mesh"], preprocess_image=False, sparse_structure_sampler_params={ "steps": min(ss_sampling_steps, 20), # L40S에서 더 많은 스텝 허용 "cfg_strength": ss_guidance_strength, }, slat_sampler_params={ "steps": min(slat_sampling_steps, 20), "cfg_strength": slat_guidance_strength, } ) # 비디오 생성 video = render_utils.render_video(outputs['gaussian'][0], num_frames=40)['color'] video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=40)['normal'] video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))] trial_id = str(uuid.uuid4()) video_path = f"{TMP_DIR}/{trial_id}.mp4" os.makedirs(os.path.dirname(video_path), exist_ok=True) imageio.mimsave(video_path, video, fps=20) state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id) if torch.cuda.is_available(): pipeline.to("cpu") return state, video_path except Exception as e: print(f"Error in image_to_3d: {str(e)}") if torch.cuda.is_available(): pipeline.to("cpu") raise e def generate_image_from_text(prompt, height, width, guidance_scale, num_steps): try: free_memory() flux_pipe = get_flux_pipe() if flux_pipe is None: raise Exception("Failed to load Flux pipeline") # L40S에 맞게 크기 제한 조정 height = min(height, 1024) width = min(width, 1024) translated_prompt = translate_if_korean(prompt) final_prompt = f"{translated_prompt}, wbgmsst, 3D, white background" with torch.cuda.amp.autocast(): output = flux_pipe( prompt=[final_prompt], height=height, width=width, guidance_scale=guidance_scale, num_inference_steps=num_steps, generator=torch.Generator(device='cuda') ) image = output.images[0] free_memory() return image except Exception as e: print(f"Error in generate_image_from_text: {str(e)}") free_memory() raise e def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[str, str]: 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 def activate_button() -> gr.Button: return gr.Button(interactive=True) def deactivate_button() -> gr.Button: return gr.Button(interactive=False) css = """ footer { visibility: hidden; } """ # Gradio 인터페이스 정의 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo: gr.Markdown(""" # Craft3D : 3D Asset Creation & Text-to-Image Generation """) with gr.Tabs(): with gr.TabItem("Image to 3D"): with gr.Row(): with gr.Column(): image_prompt = gr.Image(label="Image Prompt", image_mode="RGBA", type="pil", height=300) with gr.Accordion(label="Generation Settings", open=False): seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1) randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) gr.Markdown("Stage 1: Sparse Structure Generation") with gr.Row(): ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1) ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1) gr.Markdown("Stage 2: Structured Latent Generation") with gr.Row(): slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1) slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1) generate_btn = gr.Button("Generate") with gr.Accordion(label="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) with gr.TabItem("Text to Image"): with gr.Row(): with gr.Column(): text_prompt = gr.Textbox( label="Text Prompt", placeholder="Enter your image description...", lines=3 ) with gr.Row(): txt2img_height = gr.Slider(256, 1024, value=512, step=64, label="Height") txt2img_width = gr.Slider(256, 1024, value=512, step=64, label="Width") with gr.Row(): guidance_scale = gr.Slider(1.0, 20.0, value=7.5, label="Guidance Scale") num_steps = gr.Slider(1, 50, value=20, label="Number of Steps") generate_txt2img_btn = gr.Button("Generate Image") with gr.Column(): txt2img_output = gr.Image(label="Generated Image") trial_id = gr.Textbox(visible=False) output_buf = gr.State() # Example images with gr.Row(): examples = gr.Examples( examples=[ f'assets/example_image/{image}' for image in os.listdir("assets/example_image") ], inputs=[image_prompt], fn=preprocess_image, outputs=[trial_id, image_prompt], run_on_click=True, examples_per_page=32, # 예제 수 감소 cache_examples=False # 예제 캐싱 비활성화는 Examples 컴포넌트에서 설정 ) # Handlers image_prompt.upload( preprocess_image, inputs=[image_prompt], outputs=[trial_id, image_prompt], ) image_prompt.clear( lambda: '', outputs=[trial_id], ) generate_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], concurrency_limit=1 ).then( activate_button, outputs=[extract_glb_btn] ) extract_glb_btn.click( extract_glb, inputs=[output_buf, mesh_simplify, texture_size], outputs=[model_output, download_glb], concurrency_limit=1 ).then( activate_button, outputs=[download_glb] ) generate_txt2img_btn.click( generate_image_from_text, inputs=[text_prompt, txt2img_height, txt2img_width, guidance_scale, num_steps], outputs=[txt2img_output], concurrency_limit=1, show_progress=True # 진행 상황 표시 ) if __name__ == "__main__": import warnings warnings.filterwarnings('ignore') # CUDA 설정 확인 if torch.cuda.is_available(): print(f"Using GPU: {torch.cuda.get_device_name()}") print(f"Available GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB") # 디렉토리 생성 os.makedirs(TMP_DIR, exist_ok=True) # 메모리 정리 free_memory() # 모델 초기화 if not initialize_models(): print("Failed to initialize models") exit(1) # Gradio 앱 실행 demo.queue(max_size=2).launch( # 큐 크기 증가 share=True, max_threads=4, # 스레드 수 증가 show_error=True, server_port=7860, server_name="0.0.0.0" )