Update app.py

app.py

@@ -210,43 +210,37 @@ 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}")
         
-        # GPU로 모델 이동
-        if not g.trellis_pipeline.device.type == "cuda":
-            g.trellis_pipeline = g.trellis_pipeline.to("cuda")
-        
-        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,
-            },
-        )
+        with torch.cuda.device(0):
+            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)
-        
-        # CPU로 모델 되돌리기
-        g.trellis_pipeline = g.trellis_pipeline.to("cpu")
-        
-        return state, video_path
+            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)
+            
+            torch.cuda.empty_cache()
+            return state, video_path
         
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
-        if g.trellis_pipeline.device.type == "cuda":
-            g.trellis_pipeline = g.trellis_pipeline.to("cpu")
+        torch.cuda.empty_cache()
         return None, None
 
 
@@ -498,8 +492,7 @@ if __name__ == "__main__":
             allowed_paths=[PERSISTENT_DIR, TMP_DIR],
             server_name="0.0.0.0",
             server_port=7860,
-            show_error=True,
-            share=True
+            show_error=True
         )
         
     except Exception as e: