Update scripts/utils.py

scripts/utils.py

@@ -1,3 +1,5 @@
+import os
+import spaces
 import torch
 import numpy as np
 from PIL import Image
@@ -7,22 +9,42 @@ from pymeshlab import PercentageValue
 from pytorch3d.renderer import TexturesVertex
 from pytorch3d.structures import Meshes
 from rembg import new_session, remove
-import torch
 import torch.nn.functional as F
 from typing import List, Tuple
-from PIL import Image
 import trimesh
 
-providers = [
-    ('CUDAExecutionProvider', {
-        'device_id': 0,
-        'arena_extend_strategy': 'kSameAsRequested',
-        'gpu_mem_limit': 8 * 1024 * 1024 * 1024,
-        'cudnn_conv_algo_search': 'HEURISTIC',
-    })
-]
+# ZeroGPU 환경 감지
+IS_ZEROGPU = os.environ.get('SPACE_ID') is not None or os.environ.get('ZEROGPU') is not None
 
-session = new_session(providers=providers)
+# 전역 변수로 session 선언 (초기에는 None)
+_session = None
+_gpu_session = None
+
+def get_providers():
+    """환경에 따른 적절한 providers 반환"""
+    if IS_ZEROGPU:
+        # ZeroGPU 환경에서는 초기에 CPU만 사용
+        return ['CPUExecutionProvider']
+    else:
+        # 일반 환경에서는 CUDA 우선 사용
+        return [
+            ('CUDAExecutionProvider', {
+                'device_id': 0,
+                'arena_extend_strategy': 'kSameAsRequested',
+                'gpu_mem_limit': 8 * 1024 * 1024 * 1024,
+                'cudnn_conv_algo_search': 'HEURISTIC',
+            })
+        ]
+
+def get_session():
+    """세션을 lazy loading으로 생성"""
+    global _session
+    if _session is None:
+        _session = new_session(providers=get_providers())
+    return _session
+
+# 기존 코드와의 호환성을 위한 session 변수
+session = None  # 초기에는 None, 필요시 get_session() 사용
 
 NEG_PROMPT="sketch, sculpture, hand drawing, outline, single color, NSFW, lowres, bad anatomy,bad hands, text, error, missing fingers, yellow sleeves, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry,(worst quality:1.4),(low quality:1.4)"
 
@@ -62,7 +84,6 @@ def meshlab_mesh_to_py3dmesh(mesh: pymeshlab.Mesh) -> Meshes:
     textures = TexturesVertex(verts_features=[colors])
     return Meshes(verts=[verts], faces=[faces], textures=textures)
 
-
 def py3dmesh_to_meshlab_mesh(meshes: Meshes) -> pymeshlab.Mesh:
     colors_in = F.pad(meshes.textures.verts_features_packed().cpu().float(), [0,1], value=1).numpy().astype(np.float64)
     m1 = pymeshlab.Mesh(
@@ -72,7 +93,6 @@ def py3dmesh_to_meshlab_mesh(meshes: Meshes) -> pymeshlab.Mesh:
         v_color_matrix=colors_in)
     return m1
 
-
 def to_pyml_mesh(vertices,faces):
     m1 = pymeshlab.Mesh(
         vertex_matrix=vertices.cpu().float().numpy().astype(np.float64),
@@ -80,7 +100,6 @@ def to_pyml_mesh(vertices,faces):
     )
     return m1
 
-
 def to_py3d_mesh(vertices, faces, normals=None):
     from pytorch3d.structures import Meshes
     from pytorch3d.renderer.mesh.textures import TexturesVertex
@@ -91,7 +110,6 @@ def to_py3d_mesh(vertices, faces, normals=None):
     mesh.textures = TexturesVertex(verts_features=[normals / 2 + 0.5])
     return mesh
 
-
 def from_py3d_mesh(mesh):
     return mesh.verts_list()[0], mesh.faces_list()[0], mesh.textures.verts_features_packed()
 
@@ -126,7 +144,6 @@ def rotate_normals(normal_pils, return_types='np', rotate_direction=1) -> np.nda
             raise ValueError(f"return_types should be 'np' or 'pil', but got {return_types}")
     return ret
 
-
 def rotate_normalmap_by_angle_torch(normal_map, angle):
     """
     rotate along y-axis
@@ -140,7 +157,9 @@ def rotate_normalmap_by_angle_torch(normal_map, angle):
     return torch.matmul(normal_map.view(-1, 3), R.T).view(normal_map.shape)
 
 def do_rotate(rgba_normal, angle):
-    rgba_normal = torch.from_numpy(rgba_normal).float().cuda() / 255
+    # GPU 사용 가능 여부 확인
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    rgba_normal = torch.from_numpy(rgba_normal).float().to(device) / 255
     rotated_normal_tensor = rotate_normalmap_by_angle_torch(rgba_normal[..., :3] * 2 - 1, angle)
     rotated_normal_tensor = (rotated_normal_tensor + 1) / 2
     rotated_normal_tensor = rotated_normal_tensor * rgba_normal[:, :, [3]]    # make bg black
@@ -195,7 +214,6 @@ def change_bkgd_to_normal(normal_pils) -> List[Image.Image]:
         ret.append(Image.fromarray(rgba_normal_np.astype(np.uint8)))
     return ret
 
-
 def fix_vert_color_glb(mesh_path):
     from pygltflib import GLTF2, Material, PbrMetallicRoughness
     obj1 = GLTF2().load(mesh_path)
@@ -211,12 +229,10 @@ def fix_vert_color_glb(mesh_path):
     ))
     obj1.save(mesh_path)
 
-
 def srgb_to_linear(c_srgb):
     c_linear = np.where(c_srgb <= 0.04045, c_srgb / 12.92, ((c_srgb + 0.055) / 1.055) ** 2.4)
     return c_linear.clip(0, 1.)
 
-
 def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to_LinearRGB=True):
     # convert from pytorch3d meshes to trimesh mesh
     vertices = meshes.verts_packed().cpu().float().numpy()
@@ -239,7 +255,6 @@ def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to
         fix_vert_color_glb(save_glb_path)
     print(f"saving to {save_glb_path}")
 
-
 def save_glb_and_video(save_mesh_prefix: str, meshes: Meshes, with_timestamp=True, dist=3.5, azim_offset=180, resolution=512, fov_in_degrees=1 / 1.15, cam_type="ortho", view_padding=60, export_video=True) -> Tuple[str, str]:
     import time
     if '.' in save_mesh_prefix:
@@ -251,7 +266,6 @@ def save_glb_and_video(save_mesh_prefix: str, meshes: Meshes, with_timestamp=Tru
     save_py3dmesh_with_trimesh_fast(meshes, ret_mesh)
     return ret_mesh, None
 
-
 def simple_clean_mesh(pyml_mesh: ml.Mesh, apply_smooth=True, stepsmoothnum=1, apply_sub_divide=False, sub_divide_threshold=0.25):
     ms = ml.MeshSet()
     ms.add_mesh(pyml_mesh, "cube_mesh")
@@ -264,7 +278,6 @@ def simple_clean_mesh(pyml_mesh: ml.Mesh, apply_smooth=True, stepsmoothnum=1, ap
         ms.apply_filter("meshing_surface_subdivision_loop", iterations=2, threshold=PercentageValue(sub_divide_threshold))
     return meshlab_mesh_to_py3dmesh(ms.current_mesh())
 
-
 def expand2square(pil_img, background_color):
     width, height = pil_img.size
     if width == height:
@@ -278,13 +291,37 @@ def expand2square(pil_img, background_color):
         result.paste(pil_img, ((height - width) // 2, 0))
         return result
 
-
-def simple_preprocess(input_image, rembg_session=session, background_color=255):
+# ZeroGPU용 배경 제거 함수
+@spaces.GPU(duration=30)
+def remove_background_gpu(input_image, alpha_matting=False):
+    """GPU에서 배경 제거 실행"""
+    global _gpu_session
+    if _gpu_session is None:
+        # GPU가 할당되면 CUDA 프로바이더로 새 세션 생성
+        gpu_providers = [
+            ('CUDAExecutionProvider', {
+                'device_id': 0,
+                'arena_extend_strategy': 'kSameAsRequested',
+                'gpu_mem_limit': 8 * 1024 * 1024 * 1024,
+                'cudnn_conv_algo_search': 'HEURISTIC',
+            })
+        ]
+        _gpu_session = new_session(providers=gpu_providers)
+    return remove(input_image, alpha_matting=alpha_matting, session=_gpu_session)
+
+def simple_preprocess(input_image, rembg_session=None, background_color=255):
     RES = 2048
     input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS)
     if input_image.mode != 'RGBA':
         image_rem = input_image.convert('RGBA')
-        input_image = remove(image_rem, alpha_matting=False, session=rembg_session)
+        # ZeroGPU 환경에서는 GPU 함수 사용
+        if IS_ZEROGPU:
+            input_image = remove_background_gpu(image_rem, alpha_matting=False)
+        else:
+            # 일반 환경에서는 세션 사용
+            if rembg_session is None:
+                rembg_session = get_session()
+            input_image = remove(image_rem, alpha_matting=False, session=rembg_session)
 
     arr = np.asarray(input_image)
     alpha = np.asarray(input_image)[:, :, -1]
@@ -299,7 +336,10 @@ def simple_preprocess(input_image, rembg_session=session, background_color=255):
     input_image = expand2square(input_image, (background_color, background_color, background_color, 0))
     return input_image
 
-def init_target(img_pils, new_bkgd=(0., 0., 0.), device="cuda"):
+def init_target(img_pils, new_bkgd=(0., 0., 0.), device=None):
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    
     # Convert the background color to a PyTorch tensor
     new_bkgd = torch.tensor(new_bkgd, dtype=torch.float32).view(1, 1, 3).to(device)