Update app.py

app.py

@@ -1,18 +1,18 @@
 import gradio as gr
-from fastapi import FastAPI, UploadFile, File, Request
-from fastapi.responses import FileResponse, HTMLResponse
 import uuid
 import os
 import json
 from datetime import datetime
 from typing import Dict, List
 import shutil
-import asyncio
-from contextlib import asynccontextmanager
+import time
+import threading
 
 # Initialize data storage
 peers: Dict[str, Dict] = {}
 jobs: List[Dict] = []
+job_lock = threading.Lock()
+peer_lock = threading.Lock()
 
 # Create directories
 os.makedirs("results", exist_ok=True)
@@ -25,6 +25,7 @@ import time
 import os
 import sys
 from datetime import datetime
+import json
 
 # 설정
 PEER_ID = f"peer-{os.getenv('COMPUTERNAME', 'unknown')}-{datetime.now().strftime('%Y%m%d%H%M%S')}"
@@ -45,14 +46,48 @@ def check_gpu():
 def register_peer():
     """서버에 Peer 등록"""
     try:
-        response = requests.post(f"{SERVER_URL}/api/peers/register", params={"peer_id": PEER_ID})
+        response = requests.post(
+            f"{SERVER_URL}/run/register_peer",
+            json={"data": [PEER_ID]},
+            headers={"Content-Type": "application/json"}
+        )
         if response.status_code == 200:
-            print(f"✅ Peer 등록 성공: {PEER_ID}")
-            return True
+            result = response.json()
+            if result.get("data") and result["data"][0] == "registered":
+                print(f"✅ Peer 등록 성공: {PEER_ID}")
+                return True
     except Exception as e:
         print(f"❌ 서버 연결 실패: {e}")
     return False
 
+def send_heartbeat():
+    """서버에 heartbeat 전송"""
+    try:
+        response = requests.post(
+            f"{SERVER_URL}/run/heartbeat",
+            json={"data": [PEER_ID]},
+            headers={"Content-Type": "application/json"}
+        )
+        return response.status_code == 200
+    except:
+        return False
+
+def request_job():
+    """작업 요청"""
+    try:
+        response = requests.post(
+            f"{SERVER_URL}/run/request_job",
+            json={"data": [PEER_ID]},
+            headers={"Content-Type": "application/json"}
+        )
+        if response.status_code == 200:
+            result = response.json()
+            if result.get("data") and result["data"][0]:
+                return json.loads(result["data"][0])
+    except Exception as e:
+        print(f"작업 요청 실패: {e}")
+    return None
+
 def generate_image_cpu(prompt, output_path):
     """CPU로 테스트 이미지 생성"""
     from PIL import Image, ImageDraw, ImageFont
@@ -74,6 +109,24 @@ def generate_image_cpu(prompt, output_path):
     img.save(output_path)
     print(f"📝 테스트 이미지 생성 완료: {output_path}")
 
+def submit_result(job_id, image_path):
+    """작업 결과 제출"""
+    try:
+        # 이미지를 base64로 인코딩
+        import base64
+        with open(image_path, "rb") as f:
+            image_data = base64.b64encode(f.read()).decode()
+        
+        response = requests.post(
+            f"{SERVER_URL}/run/submit_result",
+            json={"data": [job_id, image_data]},
+            headers={"Content-Type": "application/json"}
+        )
+        return response.status_code == 200
+    except Exception as e:
+        print(f"결과 제출 실패: {e}")
+        return False
+
 def main():
     print("🚀 P2P GPU 클라이언트 시작...")
     
@@ -84,45 +137,39 @@ def main():
     while True:
         try:
             # Heartbeat
-            requests.post(f"{SERVER_URL}/api/peers/heartbeat", params={"peer_id": PEER_ID})
+            if not send_heartbeat():
+                print("서버 연결 끊김. 재등록 시도...")
+                if not register_peer():
+                    time.sleep(30)
+                    continue
             
             # 작업 요청
-            response = requests.get(f"{SERVER_URL}/api/jobs/request", params={"peer_id": PEER_ID})
-            if response.status_code == 200:
-                job_data = response.json()
+            job = request_job()
+            if job:
+                job_id = job["id"]
+                prompt = job["prompt"]
                 
-                if job_data.get("job"):
-                    job = job_data["job"]
-                    job_id = job["id"]
-                    prompt = job["prompt"]
-                    
-                    print(f"\\n📋 새 작업 수신: {prompt}")
-                    
-                    # 이미지 생성
-                    output_path = f"{job_id}.png"
-                    
-                    if check_gpu():
-                        print("🎮 GPU 모드로 생성 중...")
-                        # 실제 GPU 생성 코드가 여기 들어갑니다
-                        generate_image_cpu(prompt, output_path)
-                    else:
-                        print("💻 CPU 모드로 생성 중...")
-                        generate_image_cpu(prompt, output_path)
-                    
-                    # 결과 업로드
-                    with open(output_path, 'rb') as f:
-                        files = {'file': (output_path, f, 'image/png')}
-                        response = requests.post(
-                            f"{SERVER_URL}/api/jobs/result",
-                            params={"job_id": job_id},
-                            files=files
-                        )
-                    
-                    if response.status_code == 200:
-                        print("✅ 결과 업로드 성공")
-                    
-                    # 임시 파일 삭제
-                    os.remove(output_path)
+                print(f"\\n📋 새 작업 수신: {prompt}")
+                
+                # 이미지 생성
+                output_path = f"{job_id}.png"
+                
+                if check_gpu():
+                    print("🎮 GPU 모드로 생성 중...")
+                    # 실제 GPU 생성 코드가 여기 들어갑니다
+                    generate_image_cpu(prompt, output_path)
+                else:
+                    print("💻 CPU 모드로 생성 중...")
+                    generate_image_cpu(prompt, output_path)
+                
+                # 결과 업로드
+                if submit_result(job_id, output_path):
+                    print("✅ 결과 업로드 성공")
+                else:
+                    print("❌ 결과 업로드 실패")
+                
+                # 임시 파일 삭제
+                os.remove(output_path)
             
             time.sleep(10)  # 10초마다 확인
             
@@ -145,14 +192,15 @@ if __name__ == "__main__":
 '''
 
 # Create client files
-with open("client/peer_agent.py", "w", encoding="utf-8") as f:
-    f.write(CLIENT_CODE)
+def create_client_files():
+    with open("client/peer_agent.py", "w", encoding="utf-8") as f:
+        f.write(CLIENT_CODE)
 
-with open("client/requirements.txt", "w") as f:
-    f.write("requests\npillow\n")
+    with open("client/requirements.txt", "w") as f:
+        f.write("requests\npillow\n")
 
-with open("client/README.md", "w", encoding="utf-8") as f:
-    f.write("""# P2P GPU 클라이언트 (Windows)
+    with open("client/README.md", "w", encoding="utf-8") as f:
+        f.write("""# P2P GPU 클라이언트 (Windows)
 
 ## 설치 방법
 1. Python 3.8 이상 설치
@@ -165,141 +213,113 @@ with open("client/README.md", "w", encoding="utf-8") as f:
 - GPU가 없으면 CPU 모드로 테스트 이미지 생성
 """)
 
-# FastAPI app with lifespan
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    # Startup
-    print("P2P GPU Hub 시작...")
-    yield
-    # Shutdown
-    print("P2P GPU Hub 종료...")
-
-app = FastAPI(lifespan=lifespan)
+create_client_files()
 
-# API endpoints
-@app.get("/api/status")
-async def get_status():
-    """시스템 상태 조회"""
-    active_peers = sum(1 for p in peers.values() 
-                      if (datetime.now() - p['last_seen']).seconds < 60)
-    pending_jobs = sum(1 for j in jobs if j['status'] == 'pending')
-    completed_jobs = sum(1 for j in jobs if j['status'] == 'completed')
-    
-    recent_results = [
-        {"filename": j['filename'], "prompt": j['prompt']} 
-        for j in jobs[-10:] if j['status'] == 'completed' and 'filename' in j
-    ]
-    
-    return {
-        "active_peers": active_peers,
-        "pending_jobs": pending_jobs,
-        "completed_jobs": completed_jobs,
-        "recent_results": recent_results
-    }
-
-@app.post("/api/peers/register")
-async def register_peer(peer_id: str):
+# Cleanup old peers periodically
+def cleanup_old_peers():
+    while True:
+        time.sleep(60)  # Check every minute
+        with peer_lock:
+            current_time = datetime.now()
+            to_remove = []
+            for peer_id, peer_info in peers.items():
+                if (current_time - peer_info['last_seen']).seconds > 120:  # 2 minutes timeout
+                    to_remove.append(peer_id)
+            for peer_id in to_remove:
+                del peers[peer_id]
+
+# Start cleanup thread
+cleanup_thread = threading.Thread(target=cleanup_old_peers, daemon=True)
+cleanup_thread.start()
+
+# Gradio API functions
+def register_peer(peer_id: str):
     """Peer 등록"""
-    peers[peer_id] = {
-        "status": "idle",
-        "last_seen": datetime.now(),
-        "jobs_completed": 0
-    }
-    return {"status": "registered", "peer_id": peer_id}
-
-@app.post("/api/peers/heartbeat")
-async def heartbeat(peer_id: str):
-    """Peer 상태 업데이트"""
-    if peer_id in peers:
-        peers[peer_id]["last_seen"] = datetime.now()
-        return {"status": "alive"}
-    return {"status": "unregistered"}
-
-@app.post("/api/jobs/submit")
-async def submit_job(request: Request):
+    with peer_lock:
+        peers[peer_id] = {
+            "status": "idle",
+            "last_seen": datetime.now(),
+            "jobs_completed": 0
+        }
+    return "registered"
+
+def heartbeat(peer_id: str):
+    """Peer heartbeat"""
+    with peer_lock:
+        if peer_id in peers:
+            peers[peer_id]["last_seen"] = datetime.now()
+            return "alive"
+    return "unregistered"
+
+def submit_job(prompt: str):
     """작업 제출"""
-    data = await request.json()
+    if not prompt:
+        return "프롬프트를 입력해주세요."
+    
     job_id = str(uuid.uuid4())
-    job = {
-        "id": job_id,
-        "prompt": data.get("prompt", ""),
-        "status": "pending",
-        "created_at": datetime.now()
-    }
-    jobs.append(job)
-    return {"job_id": job_id, "status": "submitted"}
-
-@app.get("/api/jobs/request")
-async def request_job(peer_id: str):
-    """작업 요청"""
-    for job in jobs:
-        if job["status"] == "pending":
-            job["status"] = "assigned"
-            job["peer_id"] = peer_id
-            job["assigned_at"] = datetime.now()
-            return {"job": job}
+    with job_lock:
+        job = {
+            "id": job_id,
+            "prompt": prompt,
+            "status": "pending",
+            "created_at": datetime.now().isoformat()
+        }
+        jobs.append(job)
     
-    return {"job": None}
+    return f"✅ 작업이 제출되었습니다!\n작업 ID: {job_id[:8]}..."
 
-@app.post("/api/jobs/result")
-async def submit_result(job_id: str, file: UploadFile = File(...)):
+def request_job(peer_id: str):
+    """작업 요청"""
+    with job_lock:
+        for job in jobs:
+            if job["status"] == "pending":
+                job["status"] = "assigned"
+                job["peer_id"] = peer_id
+                job["assigned_at"] = datetime.now().isoformat()
+                return json.dumps(job)
+    return None
+
+def submit_result(job_id: str, image_base64: str):
     """작업 결과 제출"""
-    filename = f"{job_id}.png"
-    file_path = f"results/{filename}"
-    
-    with open(file_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
-    
-    for job in jobs:
-        if job["id"] == job_id:
-            job["status"] = "completed"
-            job["filename"] = filename
-            job["completed_at"] = datetime.now()
-            
-            if "peer_id" in job and job["peer_id"] in peers:
-                peers[job["peer_id"]]["jobs_completed"] += 1
-            break
+    import base64
     
-    return {"status": "success", "filename": filename}
-
-@app.get("/api/results/{filename}")
-async def get_result(filename: str):
-    """생성된 이미지 조회"""
-    file_path = f"results/{filename}"
-    if os.path.exists(file_path):
-        return FileResponse(file_path)
-    return {"error": "File not found"}
-
-@app.get("/api/client/{filename}")
-async def get_client_file(filename: str):
-    """클라이언트 파일 다운로드"""
-    file_path = f"client/{filename}"
-    if os.path.exists(file_path):
-        return FileResponse(file_path, filename=filename)
-    return {"error": "File not found"}
+    # base64 디코딩하여 파일로 저장
+    try:
+        image_data = base64.b64decode(image_base64)
+        filename = f"{job_id}.png"
+        file_path = f"results/{filename}"
+        
+        with open(file_path, "wb") as f:
+            f.write(image_data)
+        
+        with job_lock:
+            for job in jobs:
+                if job["id"] == job_id:
+                    job["status"] = "completed"
+                    job["filename"] = filename
+                    job["completed_at"] = datetime.now().isoformat()
+                    
+                    # Peer 통계 업데이트
+                    if "peer_id" in job:
+                        with peer_lock:
+                            if job["peer_id"] in peers:
+                                peers[job["peer_id"]]["jobs_completed"] += 1
+                    break
+        
+        return "success"
+    except Exception as e:
+        print(f"Error saving result: {e}")
+        return "error"
 
-# Gradio interface functions
-def gradio_submit_job(prompt):
-    """Gradio를 통한 작업 제출"""
-    if not prompt:
-        return "프롬프트를 입력해주세요."
+def get_status():
+    """시스템 상태 조회"""
+    with peer_lock:
+        active_peers = sum(1 for p in peers.values() 
+                          if (datetime.now() - p['last_seen']).seconds < 60)
     
-    job_id = str(uuid.uuid4())
-    job = {
-        "id": job_id,
-        "prompt": prompt,
-        "status": "pending",
-        "created_at": datetime.now()
-    }
-    jobs.append(job)
-    return f"✅ 작업이 제출되었습니다!\n작업 ID: {job_id[:8]}..."
-
-def gradio_get_status():
-    """Gradio를 통한 상태 조회"""
-    active_peers = sum(1 for p in peers.values() 
-                      if (datetime.now() - p['last_seen']).seconds < 60)
-    pending = sum(1 for j in jobs if j['status'] == 'pending')
-    completed = sum(1 for j in jobs if j['status'] == 'completed')
+    with job_lock:
+        pending = sum(1 for j in jobs if j['status'] == 'pending')
+        completed = sum(1 for j in jobs if j['status'] == 'completed')
     
     status_text = f"""## 📊 시스템 상태
 
@@ -311,28 +331,36 @@ def gradio_get_status():
 ### 최근 작업 목록
 """
     
-    # Add recent jobs
-    recent_jobs = jobs[-5:][::-1]  # Last 5 jobs, reversed
-    if recent_jobs:
-        for job in recent_jobs:
-            status_emoji = "✅" if job['status'] == 'completed' else "⏳" if job['status'] == 'pending' else "🔄"
-            status_text += f"\n{status_emoji} **{job['id'][:8]}**: {job['prompt'][:50]}... ({job['status']})"
-    else:
-        status_text += "\n아직 작업이 없습니다."
+    with job_lock:
+        recent_jobs = jobs[-5:][::-1]  # Last 5 jobs, reversed
+        if recent_jobs:
+            for job in recent_jobs:
+                status_emoji = "✅" if job['status'] == 'completed' else "⏳" if job['status'] == 'pending' else "🔄"
+                status_text += f"\n{status_emoji} **{job['id'][:8]}**: {job['prompt'][:50]}... ({job['status']})"
+        else:
+            status_text += "\n아직 작업이 없습니다."
     
     return status_text
 
-def gradio_get_gallery():
+def get_gallery():
     """갤러리용 이미지 가져오기"""
     image_files = []
-    for job in jobs[-20:][::-1]:  # Last 20 jobs, newest first
-        if job['status'] == 'completed' and 'filename' in job:
-            file_path = f"results/{job['filename']}"
-            if os.path.exists(file_path):
-                image_files.append((file_path, job['prompt']))
+    with job_lock:
+        for job in jobs[-20:][::-1]:  # Last 20 jobs, newest first
+            if job['status'] == 'completed' and 'filename' in job:
+                file_path = f"results/{job['filename']}"
+                if os.path.exists(file_path):
+                    image_files.append((file_path, job['prompt']))
     
     return image_files
 
+def download_client_file(filename):
+    """클라이언트 파일 다운로드"""
+    file_path = f"client/{filename}"
+    if os.path.exists(file_path):
+        return file_path
+    return None
+
 # Custom CSS
 custom_css = """
 .container {
@@ -386,7 +414,7 @@ with gr.Blocks(title="P2P GPU 이미지 생성 허브", css=custom_css, theme=gr
                     """)
             
             submit_btn.click(
-                fn=gradio_submit_job, 
+                fn=submit_job, 
                 inputs=prompt_input, 
                 outputs=result_text
             )
@@ -399,15 +427,14 @@ with gr.Blocks(title="P2P GPU 이미지 생성 허브", css=custom_css, theme=gr
             status_display = gr.Markdown()
             with gr.Row():
                 refresh_btn = gr.Button("🔄 새로고침", variant="secondary")
-                auto_refresh = gr.Checkbox(label="자동 새로고침 (10초)", value=True)
             
             refresh_btn.click(
-                fn=gradio_get_status, 
+                fn=get_status, 
                 outputs=status_display
             )
             
-            # Auto-refresh status
-            demo.load(fn=gradio_get_status, outputs=status_display)
+            # Auto-refresh status on load
+            demo.load(fn=get_status, outputs=status_display)
         
         with gr.Tab("🖼️ 갤러리"):
             gr.Markdown("### 최근 생성된 이미지들")
@@ -424,23 +451,42 @@ with gr.Blocks(title="P2P GPU 이미지 생성 허브", css=custom_css, theme=gr
             refresh_gallery_btn = gr.Button("🔄 갤러리 새로고침", variant="secondary")
             
             refresh_gallery_btn.click(
-                fn=gradio_get_gallery,
+                fn=get_gallery,
                 outputs=gallery
             )
             
             # Auto-load gallery
-            demo.load(fn=gradio_get_gallery, outputs=gallery)
+            demo.load(fn=get_gallery, outputs=gallery)
         
         with gr.Tab("💻 클라이언트 다운로드"):
             gr.Markdown("""
             ## Windows 클라이언트 설정 가이드
             
             ### 📥 1단계: 파일 다운로드
-            다음 파일들을 다운로드하세요:
-            - [peer_agent.py](/api/client/peer_agent.py) - 메인 클라이언트
-            - [requirements.txt](/api/client/requirements.txt) - 필요한 패키지 목록
-            - [README.md](/api/client/README.md) - 설명서
+            클라이언트 파일을 다운로드하려면 아래 버튼을 클릭하세요:
+            """)
+            
+            with gr.Row():
+                download_peer_btn = gr.Button("📥 peer_agent.py 다운로드")
+                download_req_btn = gr.Button("📥 requirements.txt 다운로드")
+                download_readme_btn = gr.Button("📥 README.md 다운로드")
+            
+            download_output = gr.File(label="다운로드 파일", visible=False)
+            
+            download_peer_btn.click(
+                fn=lambda: download_client_file("peer_agent.py"),
+                outputs=download_output
+            )
+            download_req_btn.click(
+                fn=lambda: download_client_file("requirements.txt"),
+                outputs=download_output
+            )
+            download_readme_btn.click(
+                fn=lambda: download_client_file("README.md"),
+                outputs=download_output
+            )
             
+            gr.Markdown("""
             ### 🐍 2단계: Python 설치
             - Python 3.8 이상 버전 필요
             - [python.org](https://python.org)에서 다운로드
@@ -466,19 +512,46 @@ with gr.Blocks(title="P2P GPU 이미지 생성 허브", css=custom_css, theme=gr
             - CPU 폴백 지원
             - 자동 작업 처리
             - 10초마다 새 작업 확인
-            """)
             
-            with gr.Row():
-                gr.Markdown("""
-                ### 🎮 GPU 지원
-                - NVIDIA GPU가 있으면 자동으로 활용
-                - GPU가 없어도 CPU로 테스트 가능
-                - GPU 사용률 20% 미만일 때만 작업 수행
-                """)
-
-# Mount Gradio app to FastAPI
-app = gr.mount_gradio_app(app, demo, path="/")
-
-# For local testing
+            ### 🎮 GPU 지원
+            - NVIDIA GPU가 있으면 자동으로 활용
+            - GPU가 없어도 CPU로 테스트 가능
+            - GPU 사용률 20% 미만일 때만 작업 수행
+            """)
+    
+    # Hidden API endpoints for client communication
+    gr.Interface(
+        fn=register_peer,
+        inputs=gr.Textbox(visible=False),
+        outputs=gr.Textbox(visible=False),
+        api_name="register_peer",
+        visible=False
+    )
+    
+    gr.Interface(
+        fn=heartbeat,
+        inputs=gr.Textbox(visible=False),
+        outputs=gr.Textbox(visible=False),
+        api_name="heartbeat",
+        visible=False
+    )
+    
+    gr.Interface(
+        fn=request_job,
+        inputs=gr.Textbox(visible=False),
+        outputs=gr.Textbox(visible=False),
+        api_name="request_job",
+        visible=False
+    )
+    
+    gr.Interface(
+        fn=submit_result,
+        inputs=[gr.Textbox(visible=False), gr.Textbox(visible=False)],
+        outputs=gr.Textbox(visible=False),
+        api_name="submit_result",
+        visible=False
+    )
+
+# Launch the app
 if __name__ == "__main__":
     demo.launch()