Update app.py

app.py

@@ -40,6 +40,7 @@ def initialize_models(device):
         g.trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
             "JeffreyXiang/TRELLIS-image-large"
         )
+        print("TrellisImageTo3DPipeline loaded successfully")
         
         # 이미지 생성 파이프라인
         print("Loading flux_pipe...")
@@ -48,6 +49,7 @@ def initialize_models(device):
             torch_dtype=torch.bfloat16,
             device_map="balanced"
         )
+        print("FluxPipeline loaded successfully")
         
         # Hyper-SD LoRA 로드
         print("Loading LoRA weights...")
@@ -58,6 +60,7 @@ def initialize_models(device):
         )
         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...")
@@ -77,12 +80,14 @@ 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:512"
 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'
 
 # CUDA 초기화 방지
 torch.set_grad_enabled(False)
@@ -214,35 +219,55 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         print(f"Successfully loaded image with size: {image.size}")
         
         # GPU 작업 시작
-        with torch.inference_mode():
-            outputs = g.trellis_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,
-                },
-            )
-
-            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))]
-            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)
+        with torch.cuda.device(0):
+            # 모델을 GPU로 이동
+            g.trellis_pipeline.to('cuda')
             
-            return state, video_path
+            with torch.inference_mode(), torch.cuda.amp.autocast():
+                outputs = g.trellis_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,
+                    },
+                )
+
+                # GPU에서 비디오 렌더링
+                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']
+                
+                # 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)
+                
+                # 상태 저장 전에 CPU로 이동
+                state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], new_trial_id)
+                
+                # 모델을 다시 CPU로
+                g.trellis_pipeline.to('cpu')
+                torch.cuda.empty_cache()
+                
+                return state, video_path
         
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
+        # 에러 발생 시에도 모델을 CPU로 되돌림
+        if hasattr(g.trellis_pipeline, 'to'):
+            g.trellis_pipeline.to('cpu')
+        torch.cuda.empty_cache()
         return None, None