Update app.py

app.py

@@ -1,7 +1,6 @@
 import shlex
 import subprocess
 
-
 subprocess.run(shlex.split("pip install pip==24.0"), check=True)
 subprocess.run(
     shlex.split(
@@ -14,7 +13,7 @@ subprocess.run(
     ), check=True
 )
 
-
+# 모델 체크포인트 다운로드 및 torch 설정
 if __name__ == "__main__":
     from huggingface_hub import snapshot_download
 
@@ -33,24 +32,87 @@ import gradio as gr
 from gradio_app.gradio_3dgen import create_ui as create_3d_ui
 from gradio_app.all_models import model_zoo
 
+# ===============================
+# Text-to-IMAGE 관련 API 함수 정의
+# ===============================
+def text_to_image(height, width, steps, scales, prompt, seed):
+    """
+    주어진 파라미터를 이용해 외부 API (http://211.233.58.201:7971/)의
+    /process_and_save_image 엔드포인트를 호출하여 이미지를 생성한다.
+    """
+    from gradio_client import Client
+    client = Client("http://211.233.58.201:7971/")
+    result = client.predict(
+        height,
+        width,
+        steps,
+        scales,
+        prompt,
+        seed,
+        api_name="/process_and_save_image"
+    )
+    # API 응답은 dict 형태이며, "url" 혹은 "path"에 생성된 이미지 정보가 있다.
+    # 여기서는 "url" 값을 반환하도록 한다.
+    return result.get("url", None)
 
-_TITLE = '''Text to 3D'''
+def update_random_seed():
+    """
+    외부 API의 /update_random_seed 엔드포인트를 호출하여
+    새로운 랜덤 시드 값을 가져온다.
+    """
+    from gradio_client import Client
+    client = Client("http://211.233.58.201:7971/")
+    return client.predict(api_name="/update_random_seed")
+
+# ===============================
+# UI 타이틀 및 설명
+# ===============================
+_TITLE = '''Text to 3D & Image Generation'''
 _DESCRIPTION = '''
-Text to 3D'''
+Text를 이용하여 3D 모델과 이미지를 생성할 수 있습니다.
+왼쪽 탭은 3D 모델 생성을, 오른쪽 탭은 이미지 생성을 담당합니다.
+'''
 
 def launch():
+    # 3D 모델 초기화
     model_zoo.init_models()
-        
-    with gr.Blocks(
-        title=_TITLE,
-        # theme=gr.themes.Monochrome(),
-    ) as demo:
+    
+    # Gradio Blocks 생성 (두 탭 포함)
+    with gr.Blocks(title=_TITLE) as demo:
         with gr.Row():
-            with gr.Column(scale=1):
-                gr.Markdown('# ' + _TITLE)
+            gr.Markdown('# ' + _TITLE)
         gr.Markdown(_DESCRIPTION)
-        create_3d_ui("wkl")
-
+        
+        # 탭 생성: 기존의 Text-to-3D와 새로 추가한 Text-to-IMAGE
+        with gr.Tabs():
+            with gr.Tab("Text to 3D"):
+                create_3d_ui("wkl")
+                
+            with gr.Tab("Text to IMAGE"):
+                # 이미지 생성을 위한 파라미터 입력 컴포넌트 구성
+                with gr.Row():
+                    height_slider = gr.Slider(label="Height", minimum=256, maximum=2048, step=1, value=1024)
+                    width_slider = gr.Slider(label="Width", minimum=256, maximum=2048, step=1, value=1024)
+                with gr.Row():
+                    steps_slider = gr.Slider(label="Inference Steps", minimum=1, maximum=100, step=1, value=8)
+                    scales_slider = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=3.5)
+                prompt_text = gr.Textbox(label="Image Description", placeholder="Enter prompt here", lines=2)
+                seed_number = gr.Number(label="Seed (optional, leave empty for random)", value=None)
+                
+                # 'Update Random Seed' 버튼을 누르면 API를 통해 새로운 시드값을 받아 입력란 업데이트
+                update_seed_button = gr.Button("Update Random Seed")
+                update_seed_button.click(fn=update_random_seed, inputs=[], outputs=seed_number)
+                
+                generate_button = gr.Button("Generate Image")
+                image_output = gr.Image(label="Generated Image")
+                
+                # 'Generate Image' 버튼 클릭 시 text_to_image 함수를 호출하여 결과 이미지를 출력
+                generate_button.click(
+                    fn=text_to_image,
+                    inputs=[height_slider, width_slider, steps_slider, scales_slider, prompt_text, seed_number],
+                    outputs=image_output
+                )
+                
     demo.queue().launch(share=True)
     
 if __name__ == '__main__':