Update app.py

app.py

@@ -1,17 +1,172 @@
-from fastapi import FastAPI
+import os
+import io
+import logging
+import hashlib
+from contextlib import asynccontextmanager
+from typing import Optional
+import aiohttp
 import gradio as gr
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import FileResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from PIL import Image
 
-app = FastAPI()
+# --- Configuration ---
+class Config:
+    OUTPUT_DIR = os.path.join(os.getcwd(), "comic_outputs")
+    RUNPOD_ENDPOINT = os.getenv("RUNPOD_ENDPOINT", "http://localhost:8001")
+    RUNPOD_API_KEY = os.getenv("RUNPOD_API_KEY", "")
+    MAX_RETRIES = 3
+    TIMEOUT = 30
+    GRADIO_PORT = int(os.getenv("GRADIO_PORT", 7860))
 
+# --- Logging Setup ---
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    handlers=[logging.StreamHandler()]
+)
+logger = logging.getLogger(__name__)
+
+# --- Data Models ---
+class GenerationRequest(BaseModel):
+    prompt: str
+    steps: Optional[int] = 30
+    width: Optional[int] = 768
+    height: Optional[int] = 768
+    seed: Optional[int] = None
+
+# --- Application Setup ---
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Modern lifespan handler with resource management"""
+    # Startup
+    os.makedirs(Config.OUTPUT_DIR, exist_ok=True)
+    logger.info("🚀 Starting AI Comic Maker API")
+    logger.info(f"📁 Output directory: {Config.OUTPUT_DIR}")
+    
+    # Create aiohttp session
+    app.state.http_session = aiohttp.ClientSession()
+    
+    yield
+    
+    # Shutdown
+    await app.state.http_session.close()
+    logger.info("🛑 Graceful shutdown complete")
+
+app = FastAPI(
+    title="AI Comic Maker API",
+    description="Generate comic art using Stable Diffusion",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# CORS Middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# --- Helper Functions ---
+def sanitize_filename(prompt: str) -> str:
+    """Create safe filename from prompt"""
+    prompt_hash = hashlib.md5(prompt.encode()).hexdigest()
+    return f"comic_{prompt_hash}.png"
+
+async def generate_image(payload: dict) -> bytes:
+    """Call RunPod API with retry logic"""
+    headers = {"Authorization": f"Bearer {Config.RUNPOD_API_KEY}"} if Config.RUNPOD_API_KEY else {}
+    
+    for attempt in range(Config.MAX_RETRIES):
+        try:
+            async with app.state.http_session.post(
+                f"{Config.RUNPOD_ENDPOINT}/generate",
+                headers=headers,
+                json={"input": payload},
+                timeout=Config.TIMEOUT
+            ) as response:
+                response.raise_for_status()
+                return await response.read()
+        except Exception as e:
+            if attempt == Config.MAX_RETRIES - 1:
+                raise
+            logger.warning(f"Retry {attempt + 1} for generation request")
+            await asyncio.sleep(1)
+
+# --- API Endpoints ---
 @app.get("/health")
-def health_check():
-    return {"status": "ok"}
+async def health_check():
+    """Liveness probe endpoint"""
+    return JSONResponse(content={"status": "healthy"})
+
+@app.post("/generate")
+async def generate_comic(request: GenerationRequest):
+    """Generate comic image from text prompt"""
+    try:
+        logger.info(f"🎨 Generating comic: '{request.prompt[:50]}...'")
+        
+        # Call RunPod API
+        image_data = await generate_image({
+            "prompt": request.prompt,
+            "steps": request.steps,
+            "width": request.width,
+            "height": request.height,
+            "seed": request.seed
+        })
+        
+        # Save image
+        output_filename = sanitize_filename(request.prompt)
+        output_path = os.path.join(Config.OUTPUT_DIR, output_filename)
+        
+        with Image.open(io.BytesIO(image_data)) as img:
+            img.save(output_path)
+        
+        logger.info(f"✅ Saved comic to {output_path}")
+        return FileResponse(output_path, media_type="image/png")
+
+    except Exception as e:
+        logger.error(f"🔥 Generation failed: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# --- Gradio Interface ---
+def gradio_generate(prompt: str, steps: int = 30):
+    """Gradio wrapper for generation"""
+    try:
+        response = requests.post(
+            f"http://localhost:{Config.GRADIO_PORT}/generate",
+            json={"prompt": prompt, "steps": steps}
+        )
+        response.raise_for_status()
+        return response.content
+    except Exception as e:
+        return str(e)
 
-# Basic Gradio interface
-demo = gr.Interface(lambda x: x, "text", "text")
+gradio_interface = gr.Interface(
+    fn=gradio_generate,
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Slider(10, 100, value=30, label="Steps")
+    ],
+    outputs=gr.Image(label="Generated Comic"),
+    title="AI Comic Maker",
+    description="Generate comic art from text prompts"
+)
 
-app = gr.mount_gradio_app(app, demo, path="/")
+# --- Mount Gradio to FastAPI ---
+app = gr.mount_gradio_app(app, gradio_interface, path="/gradio")
 
+# --- Main Entry Point ---
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    import asyncio
+    
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=Config.GRADIO_PORT,
+        log_level="info",
+        reload=False
+    )