Update app.py

app.py

@@ -1,191 +1,17 @@
-#!/usr/bin/env python3
-import os
-import io
-import logging
-import hashlib
-from contextlib import asynccontextmanager
-from typing import Optional
-from fastapi import FastAPI, HTTPException, status
-from fastapi.responses import FileResponse, JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, Field
-from PIL import Image
-import requests
-import aiohttp
+from fastapi import FastAPI
+import gradio as gr
 
-# --- 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
+app = FastAPI()
 
-# --- Logging Setup ---
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-    handlers=[
-        logging.StreamHandler(),
-        logging.FileHandler("comic_maker.log")
-    ]
-)
-logger = logging.getLogger(__name__)
+@app.get("/health")
+def health_check():
+    return {"status": "ok"}
 
-# --- Data Models ---
-class GenerationRequest(BaseModel):
-    prompt: str = Field(..., min_length=3, max_length=500)
-    steps: Optional[int] = Field(30, ge=10, le=100)
-    width: Optional[int] = Field(768, ge=256, le=1024)
-    height: Optional[int] = Field(768, ge=256, le=1024)
-    seed: Optional[int] = Field(None, ge=0)
+# Basic Gradio interface
+demo = gr.Interface(lambda x: x, "text", "text")
 
-# --- 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(f"🚀 Starting AI Comic Maker API")
-    logger.info(f"📁 Output directory: {Config.OUTPUT_DIR}")
-    logger.info(f"🔗 RunPod endpoint: {Config.RUNPOD_ENDPOINT}")
-    
-    # Create aiohttp session for async requests
-    app.state.http_session = aiohttp.ClientSession()
-    
-    yield
-    
-    # Shutdown
-    await app.state.http_session.close()
-    logger.info("🛑 Graceful shutdown complete")
+app = gr.mount_gradio_app(app, demo, path="/")
 
-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", tags=["monitoring"])
-async def health_check():
-    """Service health endpoint"""
-    return JSONResponse(
-        content={
-            "status": "healthy",
-            "version": app.version,
-            "environment": os.getenv("ENVIRONMENT", "development")
-        }
-    )
-
-@app.post(
-    "/generate",
-    response_class=FileResponse,
-    status_code=status.HTTP_200_OK,
-    tags=["generation"]
-)
-async def generate_comic(request: GenerationRequest):
-    """
-    Generate comic art from text prompt
-    
-    - **prompt**: Description of desired comic (3-500 chars)
-    - **steps**: Generation steps (10-100, default 30)
-    - **width**: Image width (256-1024, default 768)
-    - **height**: Image height (256-1024, default 768)
-    - **seed**: Optional random seed
-    """
-    try:
-        logger.info(f"🎨 Generating comic: '{request.prompt[:50]}...'")
-        
-        # Prepare payload
-        payload = {
-            "prompt": request.prompt,
-            "steps": request.steps,
-            "width": request.width,
-            "height": request.height,
-            "seed": request.seed
-        }
-        
-        # Generate image
-        image_data = await generate_image(payload)
-        
-        # Save and return
-        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",
-            filename=output_filename
-        )
-        
-    except aiohttp.ClientError as e:
-        logger.error(f"🚨 Network error: {str(e)}")
-        raise HTTPException(
-            status_code=status.HTTP_502_BAD_GATEWAY,
-            detail="Generation service unavailable"
-        )
-    except Exception as e:
-        logger.error(f"🔥 Generation failed: {str(e)}")
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail="Comic generation failed"
-        )
-
-# --- Main Entry Point ---
-if __name__ == "__main__":
-    import uvicorn
-    import asyncio
-    
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=int(os.getenv("PORT", "8000")),
-        log_level="info",
-        reload=False,
-        timeout_keep_alive=30
-    )
-
-# At the bottom of app.py
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))
+    uvicorn.run(app, host="0.0.0.0", port=7860)