Update app.py

app.py

@@ -1,58 +1,410 @@
 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
-import random
+from typing import Dict, List
+import shutil
+import asyncio
+from contextlib import asynccontextmanager
 
-# 간단한 작업 목록
-jobs = []
+# Initialize data storage
+peers: Dict[str, Dict] = {}
+jobs: List[Dict] = []
 
-def submit_job(prompt):
-    """작업 제출"""
+# Create directories
+os.makedirs("results", exist_ok=True)
+os.makedirs("client", exist_ok=True)
+
+# Client code
+CLIENT_CODE = '''import requests
+import subprocess
+import time
+import os
+import sys
+from datetime import datetime
+
+# Configuration
+PEER_ID = f"peer-{os.getenv('COMPUTERNAME', 'unknown')}-{datetime.now().strftime('%Y%m%d%H%M%S')}"
+SERVER_URL = "https://your-username-your-space.hf.space"  # Replace with actual Space URL
+
+def check_gpu():
+    """Check GPU availability"""
+    try:
+        result = subprocess.run(['nvidia-smi', '--query-gpu=utilization.gpu', '--format=csv,noheader,nounits'], 
+                              capture_output=True, text=True)
+        if result.returncode == 0:
+            gpu_usage = int(result.stdout.strip())
+            return gpu_usage < 20  # GPU is idle if usage < 20%
+    except:
+        print("GPU not found. Running in CPU mode.")
+        return False
+
+def register_peer():
+    """Register peer with server"""
+    try:
+        response = requests.post(f"{SERVER_URL}/api/peers/register", params={"peer_id": PEER_ID})
+        if response.status_code == 200:
+            print(f"✅ Peer registered: {PEER_ID}")
+            return True
+    except Exception as e:
+        print(f"❌ Server connection failed: {e}")
+    return False
+
+def generate_image_cpu(prompt, output_path):
+    """Generate test image using CPU"""
+    from PIL import Image, ImageDraw, ImageFont
+    
+    img = Image.new('RGB', (512, 512), color='white')
+    draw = ImageDraw.Draw(img)
+    
+    # Draw prompt text
+    text = f"Prompt: {prompt[:50]}..."
+    draw.text((10, 10), text, fill='black')
+    draw.text((10, 40), f"Generated by: {PEER_ID}", fill='gray')
+    draw.text((10, 70), f"Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", fill='gray')
+    
+    img.save(output_path)
+    print(f"📝 Test image generated: {output_path}")
+
+def main():
+    print("🚀 Starting P2P GPU Client...")
+    
+    if not register_peer():
+        print("Server registration failed. Exiting.")
+        return
+    
+    while True:
+        try:
+            # Heartbeat
+            requests.post(f"{SERVER_URL}/api/peers/heartbeat", params={"peer_id": PEER_ID})
+            
+            # Request job
+            response = requests.get(f"{SERVER_URL}/api/jobs/request", params={"peer_id": PEER_ID})
+            if response.status_code == 200:
+                job_data = response.json()
+                
+                if job_data.get("job"):
+                    job = job_data["job"]
+                    job_id = job["id"]
+                    prompt = job["prompt"]
+                    
+                    print(f"\\n📋 New job received: {prompt}")
+                    
+                    # Generate image
+                    output_path = f"{job_id}.png"
+                    
+                    if check_gpu():
+                        print("🎮 Generating with GPU...")
+                        # Actual GPU generation code would go here
+                        generate_image_cpu(prompt, output_path)
+                    else:
+                        print("💻 Generating with CPU...")
+                        generate_image_cpu(prompt, output_path)
+                    
+                    # Upload result
+                    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("✅ Result uploaded successfully")
+                    
+                    # Clean up
+                    os.remove(output_path)
+            
+            time.sleep(10)  # Check every 10 seconds
+            
+        except KeyboardInterrupt:
+            print("\\n👋 Shutting down")
+            break
+        except Exception as e:
+            print(f"⚠️ Error: {e}")
+            time.sleep(30)
+
+if __name__ == "__main__":
+    # Check required packages
+    try:
+        import PIL
+    except ImportError:
+        print("Installing required packages...")
+        subprocess.run([sys.executable, "-m", "pip", "install", "pillow", "requests"])
+    
+    main()
+'''
+
+# 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/README.md", "w", encoding="utf-8") as f:
+    f.write("""# P2P GPU Client for Windows
+
+## Installation
+1. Install Python 3.8+
+2. Run `pip install -r requirements.txt`
+3. Update SERVER_URL in `peer_agent.py` with actual Hugging Face Space URL
+4. Run `python peer_agent.py`
+
+## GPU Support
+- Automatically detects NVIDIA GPU if available
+- Falls back to CPU mode for testing
+""")
+
+# FastAPI app with lifespan
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup
+    print("Starting P2P GPU Hub...")
+    yield
+    # Shutdown
+    print("Shutting down P2P GPU Hub...")
+
+app = FastAPI(lifespan=lifespan)
+
+# API endpoints
+@app.get("/api/status")
+async def get_status():
+    """Get system 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):
+    """Register a 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):
+    """Update peer status"""
+    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):
+    """Submit a job"""
+    data = await request.json()
+    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):
+    """Request a job for processing"""
+    for job in jobs:
+        if job["status"] == "pending":
+            job["status"] = "assigned"
+            job["peer_id"] = peer_id
+            job["assigned_at"] = datetime.now()
+            return {"job": job}
+    
+    return {"job": None}
+
+@app.post("/api/jobs/result")
+async def submit_result(job_id: str, file: UploadFile = File(...)):
+    """Submit job result"""
+    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
+    
+    return {"status": "success", "filename": filename}
+
+@app.get("/api/results/{filename}")
+async def get_result(filename: str):
+    """Get generated image"""
+    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):
+    """Download client file"""
+    file_path = f"client/{filename}"
+    if os.path.exists(file_path):
+        return FileResponse(file_path, filename=filename)
+    return {"error": "File not found"}
+
+# Gradio interface functions
+def gradio_submit_job(prompt):
+    """Submit job through Gradio"""
     if not prompt:
-        return "프롬프트를 입력해주세요."
+        return "Please enter a prompt"
     
-    job_id = f"JOB-{random.randint(1000, 9999)}"
-    jobs.append({
+    job_id = str(uuid.uuid4())
+    job = {
         "id": job_id,
         "prompt": prompt,
-        "time": datetime.now().strftime("%H:%M:%S")
-    })
-    
-    return f"✅ 작업 제출됨: {job_id}\n프롬프트: {prompt}"
-
-def get_status():
-    """상태 확인"""
-    if not jobs:
-        return "아직 작업이 없습니다."
-    
-    status = f"전체 작업: {len(jobs)}개\n\n최근 작업:\n"
-    for job in jobs[-5:]:
-        status += f"- {job['id']}: {job['prompt'][:30]}... ({job['time']})\n"
-    
-    return status
-
-# Gradio 인터페이스
-demo = gr.Interface(
-    fn=submit_job,
-    inputs=gr.Textbox(label="이미지 프롬프트", placeholder="원하는 이미지를 설명하세요..."),
-    outputs=gr.Textbox(label="결과"),
-    title="P2P GPU 이미지 생��� 허브",
-    description="분산 GPU를 활용한 이미지 생성 시스템"
-)
-
-# 상태 확인용 두 번째 인터페이스
-status_demo = gr.Interface(
-    fn=get_status,
-    inputs=None,
-    outputs=gr.Textbox(label="시스템 상태"),
-    title="시스템 상태"
-)
-
-# 탭으로 결합
-app = gr.TabbedInterface(
-    [demo, status_demo],
-    ["작업 제출", "시스템 상태"]
-)
+        "status": "pending",
+        "created_at": datetime.now()
+    }
+    jobs.append(job)
+    return f"Job submitted successfully! Job ID: {job_id}"
+
+def gradio_get_status():
+    """Get status through 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')
+    
+    status_text = f"""### System Status
+- Active Peers: {active_peers}
+- Pending Jobs: {pending}
+- Completed Jobs: {completed}
+
+### Recent Jobs
+"""
+    
+    # Add recent jobs
+    recent_jobs = jobs[-5:][::-1]  # Last 5 jobs, reversed
+    for job in recent_jobs:
+        status_text += f"\n- **{job['id'][:8]}...**: {job['prompt'][:50]}... ({job['status']})"
+    
+    return status_text
+
+def gradio_get_gallery():
+    """Get completed images for gallery"""
+    image_files = []
+    for job in jobs[-20:]:  # Last 20 jobs
+        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
+
+# Create Gradio interface
+with gr.Blocks(title="P2P GPU Image Generation Hub") as demo:
+    gr.Markdown("# 🤖 P2P GPU Image Generation Hub")
+    gr.Markdown("Distributed image generation using idle GPUs from peer nodes")
+    
+    with gr.Tabs():
+        with gr.Tab("Submit Job"):
+            with gr.Row():
+                with gr.Column():
+                    prompt_input = gr.Textbox(
+                        label="Image Prompt", 
+                        placeholder="Describe the image you want to generate...",
+                        lines=3
+                    )
+                    submit_btn = gr.Button("Submit Job", variant="primary")
+                    result_text = gr.Textbox(label="Result", interactive=False)
+                    
+                    submit_btn.click(
+                        fn=gradio_submit_job, 
+                        inputs=prompt_input, 
+                        outputs=result_text
+                    )
+        
+        with gr.Tab("System Status"):
+            status_display = gr.Markdown()
+            refresh_btn = gr.Button("Refresh Status")
+            
+            refresh_btn.click(
+                fn=gradio_get_status, 
+                outputs=status_display
+            )
+            
+            # Auto-refresh status on load
+            demo.load(fn=gradio_get_status, outputs=status_display)
+        
+        with gr.Tab("Gallery"):
+            gallery = gr.Gallery(
+                label="Generated Images",
+                show_label=True,
+                elem_id="gallery",
+                columns=3,
+                rows=2,
+                height="auto"
+            )
+            refresh_gallery_btn = gr.Button("Refresh Gallery")
+            
+            refresh_gallery_btn.click(
+                fn=gradio_get_gallery,
+                outputs=gallery
+            )
+            
+            # Auto-load gallery on tab load
+            demo.load(fn=gradio_get_gallery, outputs=gallery)
+        
+        with gr.Tab("Download Client"):
+            gr.Markdown("""
+            ## Windows Client Setup
+            
+            1. Download the client files:
+               - [peer_agent.py](/api/client/peer_agent.py)
+               - [requirements.txt](/api/client/requirements.txt)
+               - [README.md](/api/client/README.md)
+            
+            2. Install Python 3.8 or higher
+            
+            3. Install requirements:
+               ```bash
+               pip install -r requirements.txt
+               ```
+            
+            4. Update the SERVER_URL in peer_agent.py with this Space's URL
+            
+            5. Run the client:
+               ```bash
+               python peer_agent.py
+               ```
+            
+            The client will automatically detect GPU availability and start processing jobs.
+            """)
+
+# Mount Gradio app to FastAPI
+app = gr.mount_gradio_app(app, demo, path="/")
 
+# For Hugging Face Spaces
 if __name__ == "__main__":
-    app.launch(server_name="0.0.0.0", server_port=7860)
-    
+    demo.launch()