Update app.py

app.py

@@ -38,8 +38,7 @@ def initialize_models(device):
         print("Initializing models...")
         # 3D 생성 파이프라인
         g.trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large",
-            torch_dtype=torch.float32  # 명시적으로 dtype 지정
+            "JeffreyXiang/TRELLIS-image-large"
         )
         
         # 이미지 생성 파이프라인
@@ -78,15 +77,12 @@ 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['TORCH_USE_CUDA_DSA'] = '1'
-os.environ['CUDA_VISIBLE_DEVICES'] = '0'
-os.environ['XFORMERS_FORCE_DISABLE_TRITON'] = '1'  # xformers 관련 설정 추가
+os.environ['XFORMERS_FORCE_DISABLE_TRITON'] = '1'
 
 # CUDA 초기화 방지
 torch.set_grad_enabled(False)
@@ -217,47 +213,36 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         image = Image.open(image_path)
         print(f"Successfully loaded image with size: {image.size}")
         
-        # Move model to GPU
-        g.trellis_pipeline.to("cuda")
-        
-        try:
-            with torch.inference_mode():
-                with 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 작업 시작
+        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)
-                    
-                    return state, video_path
-        finally:
-            # Move model back to CPU
-            g.trellis_pipeline.to("cpu")
-            torch.cuda.empty_cache()
+            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)
+            
+            return state, video_path
         
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
-        if hasattr(g.trellis_pipeline, 'device') and g.trellis_pipeline.device.type == "cuda":
-            g.trellis_pipeline.to("cpu")
-        torch.cuda.empty_cache()
         return None, None