Update app.py

app.py

@@ -1,7 +1,11 @@
+import spaces
 import shlex
 import subprocess
 import os
 
+# ZeroGPU 환경 설정
+os.environ['CUDA_VISIBLE_DEVICES'] = ''  # 초기화 시점에는 GPU 비활성화
+
 subprocess.run(shlex.split("pip install pip==24.0"), check=True)
 subprocess.run(
     shlex.split(
@@ -36,12 +40,16 @@ from gradio_app.all_models import model_zoo
 # ===============================
 # Text-to-IMAGE 관련 API 함수 정의
 # ===============================
+@spaces.GPU(duration=60)  # GPU 사용 시간 60초로 설정
 def text_to_image(height, width, steps, scales, prompt, seed):
     """
     주어진 파라미터를 이용해 외부 API의 /process_and_save_image 엔드포인트를 호출하여 이미지를 생성한다.
     """
+    # GPU가 할당된 상태에서 실행
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    
     from gradio_client import Client
-    client = Client(os.getenv("CLIENT_API"))  # 기본값 설정
+    client = Client(os.getenv("CLIENT_API"))
     result = client.predict(
         height,
         width,
@@ -61,9 +69,17 @@ def update_random_seed():
     외부 API의 /update_random_seed 엔드포인트를 호출하여 새로운 랜덤 시드 값을 가져온다.
     """
     from gradio_client import Client
-    client = Client(os.getenv("CLIENT_API"))  # 기본값 설정
+    client = Client(os.getenv("CLIENT_API"))
     return client.predict(api_name="/update_random_seed")
 
+# 3D 생성 함수를 위한 래퍼 (GPU 데코레이터 적용)
+@spaces.GPU(duration=120)  # 3D 생성은 더 많은 시간 필요
+def generate_3d_wrapper(*args, **kwargs):
+    """3D 생성 함수를 GPU 환경에서 실행"""
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    # 실제 3D 생성 로직이 여기서 실행됨
+    # model_zoo의 함수들이 여기서 호출될 것임
+    return model_zoo.generate_3d(*args, **kwargs)
 
 _TITLE = '''✨ 3D LLAMA Studio'''
 _DESCRIPTION = '''
@@ -74,6 +90,8 @@ This platform offers two powerful features:
 2. **Text to Styled Image**: Create artistic images that can be used for 3D generation
 
 *Note: Both English and Korean prompts are supported (영어와 한글 프롬프트 모두 지원됩니다)*
+
+**Running on ZeroGPU** 🚀
 '''
 
 # CSS 스타일 밝은 테마로 수정
@@ -129,10 +147,20 @@ custom_css = """
     border: 1px solid #e0e0e0 !important;
     background-color: #ffffff !important;
 }
+.zerogpu-badge {
+    background-color: #4CAF50;
+    color: white;
+    padding: 5px 10px;
+    border-radius: 5px;
+    font-size: 14px;
+    margin-left: 10px;
+}
 """
 
 # Gradio 테마 설정 수정
 def launch():
+    # CPU 모드로 모델 초기화
+    os.environ['CUDA_VISIBLE_DEVICES'] = ''
     model_zoo.init_models()
     
     with gr.Blocks(
@@ -147,7 +175,10 @@ def launch():
     ) as demo:
 
         with gr.Row():
-            gr.Markdown('# ' + _TITLE, elem_classes="main-title")
+            with gr.Column():
+                gr.Markdown('# ' + _TITLE, elem_classes="main-title")
+            with gr.Column():
+                gr.HTML('<span class="zerogpu-badge">ZeroGPU Enabled</span>')
         gr.Markdown(_DESCRIPTION)
         
         with gr.Tabs() as tabs: