Update app.py

app.py

@@ -206,6 +206,7 @@ TRANSLATOR = None
 # API URLs
 TEXT2IMG_API_URL = "http://211.233.58.201:7896"
 VIDEO_API_URL = "http://211.233.58.201:7875"
+AVATAR_API_URL = "http://211.233.58.201:7788"
 
 # Image size presets
 IMAGE_PRESETS = {
@@ -987,6 +988,69 @@ def merge_videos_with_audio(video_files, audio_file, audio_mode, audio_volume, o
         traceback.print_exc()
         return None, f"❌ 오류 발생: {str(e)}"
 
+@spaces.GPU(duration=180)
+def generate_avatar_animation(portrait_image, driving_audio, guidance_scale, inference_steps):
+    """이미지와 오디오로 아바타 애니메이션 생성"""
+    if portrait_image is None:
+        return None, None, "초상화 이미지를 업로드해주세요."
+    
+    if driving_audio is None:
+        return None, None, "오디오 파일을 업로드해주세요."
+    
+    try:
+        # 상태 업데이트
+        status = "아바타 생성 시작..."
+        
+        # 이미지 저장
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as fp:
+            portrait_path = fp.name
+            if isinstance(portrait_image, np.ndarray):
+                Image.fromarray(portrait_image).save(portrait_path)
+            else:
+                Image.open(portrait_image).save(portrait_path)
+        
+        # 오디오 경로 확인
+        if isinstance(driving_audio, str):
+            audio_path = driving_audio
+        else:
+            audio_path = driving_audio
+        
+        logging.info(f"Portrait: {portrait_path}")
+        logging.info(f"Audio: {audio_path}")
+        logging.info(f"Guidance Scale: {guidance_scale}")
+        logging.info(f"Inference Steps: {inference_steps}")
+        
+        # API 호출
+        client = Client(AVATAR_API_URL)
+        result = client.predict(
+            handle_file(portrait_path),
+            handle_file(audio_path),
+            float(guidance_scale),
+            float(inference_steps),
+            api_name="/generate_animation"
+        )
+        
+        # 임시 파일 삭제
+        os.unlink(portrait_path)
+        
+        if result and len(result) >= 2:
+            animation_result = result[0]
+            comparison_result = result[1]
+            
+            # 비디오 경로 추출
+            animation_video = animation_result.get("video") if isinstance(animation_result, dict) else None
+            comparison_video = comparison_result.get("video") if isinstance(comparison_result, dict) else None
+            
+            return animation_video, comparison_video, "✅ 아바타 애니메이션 생성 완료!"
+        else:
+            return None, None, "❌ API 응답이 예상과 다릅니다."
+            
+    except Exception as e:
+        logging.error(f"Avatar generation error: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return None, None, f"❌ 오류 발생: {str(e)}"
+
 # CSS
 css = """
 :root {
@@ -1012,7 +1076,7 @@ css = """
     padding: 20px !important;
     margin-bottom: 20px !important;
 }
-#generate-btn, #video-btn, #outpaint-btn, #preview-btn, #audio-btn, #bg-remove-btn, #merge-btn {
+#generate-btn, #video-btn, #outpaint-btn, #preview-btn, #audio-btn, #bg-remove-btn, #merge-btn, #avatar-btn {
     background: linear-gradient(135deg, #ff9a9e, #fad0c4) !important;
     font-size: 1.1rem !important;
     padding: 12px 24px !important;
@@ -1419,6 +1483,73 @@ with demo:
                         긴 비디오는 작은 조각으로 나누어 처리하세요.
                         """)
     
+        # 여섯 번째 탭: 이미지to아바타
+        with gr.Tab("이미지to아바타", elem_classes="tabitem"):
+            with gr.Row(equal_height=True):
+                # 입력 컬럼
+                with gr.Column(scale=1):
+                    with gr.Group(elem_classes="panel-box"):
+                        gr.Markdown("### 🖼️ 초상화 이미지 업로드")
+                        
+                        avatar_portrait = gr.Image(
+                            label="초상화 이미지 (아무 비율 가능)",
+                            type="filepath"
+                        )
+                        
+                    with gr.Group(elem_classes="panel-box"):
+                        gr.Markdown("### 🎵 오디오 업로드")
+                        
+                        avatar_audio = gr.Audio(
+                            label="구동 오디오",
+                            type="filepath",
+                            sources=["upload"]
+                        )
+                        
+                    with gr.Group(elem_classes="panel-box"):
+                        gr.Markdown("### ⚙️ 생성 설정")
+                        
+                        avatar_guidance = gr.Slider(
+                            minimum=1.0,
+                            maximum=7.0,
+                            value=2.5,
+                            step=0.1,
+                            label="가이던스 스케일",
+                            info="값이 높을수록 더 강한 ��이드"
+                        )
+                        
+                        avatar_steps = gr.Slider(
+                            minimum=5,
+                            maximum=30,
+                            value=15,
+                            step=1,
+                            label="추론 스텝",
+                            info="스텝이 많을수록 품질 향상"
+                        )
+                        
+                        avatar_btn = gr.Button("🎭 아바타 생성", variant="primary", elem_id="avatar-btn")
+                
+                # 출력 컬럼
+                with gr.Column(scale=1):
+                    with gr.Group(elem_classes="panel-box"):
+                        gr.Markdown("### 🎬 생성 결과")
+                        
+                        avatar_status = gr.Textbox(label="처리 상태", interactive=False)
+                        avatar_result = gr.Video(label="애니메이션 결과")
+                        avatar_comparison = gr.Video(label="원본-애니메이션 비교")
+                        
+                        gr.Markdown("""
+                        ### ℹ️ 사용 방법
+                        1. 정면을 보고 있는 초상화 이미지를 업로드하세요
+                        2. 음성이나 음악 오디오 파일을 업로드하세요
+                        3. 가이던스와 스텝 설정을 조정하세요
+                        4. '아바타 생성' 버튼을 클릭하세요
+                        
+                        **팁**:
+                        - 정면 얼굴 이미지가 가장 좋은 결과를 만듭니다
+                        - 오디오의 음성이 명확할수록 립싱크가 정확합니다
+                        - 가이던스를 높이면 움직임이 강해집니다
+                        """)
+    
     # 모델 로드 함수 실행
     def on_demo_load():
         try:
@@ -1508,6 +1639,13 @@ with demo:
         outputs=[stream_image, output_bg_video, time_textbox]
     )
     
+    # 이벤트 연결 - 여섯 번째 탭 (이미지to아바타)
+    avatar_btn.click(
+        generate_avatar_animation,
+        inputs=[avatar_portrait, avatar_audio, avatar_guidance, avatar_steps],
+        outputs=[avatar_result, avatar_comparison, avatar_status]
+    )
+    
     # 데모 로드 시 실행
     demo.load(on_demo_load, outputs=model_status)