Create MEDIA_CLICKS_APP.PY

MEDIA_CLICKS_APP.PY

@@ -0,0 +1,548 @@
+import os
+import io
+import json
+import traceback
+from datetime import datetime,timedelta
+from typing import Optional
+import time
+import uuid
+import boto3
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request, Depends
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from pymongo import MongoClient
+import gridfs
+from gridfs.errors import NoFile
+from bson.objectid import ObjectId
+from PIL import Image
+from fastapi.concurrency import run_in_threadpool
+import shutil
+import firebase_admin
+from firebase_admin import credentials, auth
+from PIL import Image
+from huggingface_hub import InferenceClient
+# ---------------------------------------------------------------------
+# Load Firebase Config from env (stringified JSON)
+# ---------------------------------------------------------------------
+firebase_config_json = os.getenv("firebase_config")
+if not firebase_config_json:
+    raise RuntimeError("❌ Missing Firebase config in environment variable 'firebase_config'")
+
+try:
+    firebase_creds_dict = json.loads(firebase_config_json)
+    cred = credentials.Certificate(firebase_creds_dict)
+    firebase_admin.initialize_app(cred)
+except Exception as e:
+    raise RuntimeError(f"Failed to initialize Firebase Admin SDK: {e}")
+
+# ---------------------------------------------------------------------
+# Hugging Face setup
+# ---------------------------------------------------------------------
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise RuntimeError("HF_TOKEN not set in environment variables")
+
+hf_client = InferenceClient(token=HF_TOKEN)
+# ---------------------------------------------------------------------
+# MODEL SELECTION
+# ---------------------------------------------------------------------
+genai = None  # ✅ IMPORTANT: module-level declaration
+MODEL = os.getenv("IMAGE_MODEL", "GEMINI").upper()
+GEMINI_FORCE_CATEGORY_ID = "69368e741224bcb6bdb98076"
+# ---------------------------------------------------------------------
+# Gemini setup (ONLY used if MODEL == "GEMINI")
+# ---------------------------------------------------------------------
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+GEMINI_IMAGE_MODEL = os.getenv("GEMINI_IMAGE_MODEL", "gemini-2.5-flash-image")
+
+# ---------------------------------------------------------------------
+# MongoDB setup
+# ---------------------------------------------------------------------
+# MONGODB_URI=os.getenv("MONGODB_URI")
+# DB_NAME = "polaroid_db"
+
+# mongo = MongoClient(MONGODB_URI)
+# db = mongo[DB_NAME]
+# fs = gridfs.GridFS(db)
+# logs_collection = db["logs"]
+# ---------------- Logs MongoDB (NEW) ----------------
+MONGODB_URI = os.getenv("MONGODB_URI")
+POLAROID_LOGS_MONGO_URL = os.getenv("POLAROID_LOGS_MONGO_URL")
+
+DB_NAME = "polaroid_db"
+
+# Main DB (images, gridfs)
+mongo = MongoClient(MONGODB_URI)
+db = mongo[DB_NAME]
+fs = gridfs.GridFS(db)
+
+# Separate Logs DB
+logs_client = None
+logs_collection = None
+
+if POLAROID_LOGS_MONGO_URL:
+    logs_client = MongoClient(POLAROID_LOGS_MONGO_URL)
+    logs_db = logs_client["logs"]          # FORCE DB NAME
+    logs_collection = logs_db["polaroid"]  # FORCE COLLECTION
+else:
+    raise RuntimeError("POLAROID_LOGS_MONGO_URL not set")
+# ---------------------------------------------------------------------
+# DigitalOcean Spaces setup
+# ---------------------------------------------------------------------
+DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")
+DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")
+DO_SPACES_REGION = os.getenv("DO_SPACES_REGION", "blr1")
+DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")
+DO_SPACES_ENDPOINT = os.getenv("DO_SPACES_ENDPOINT", f"https://{DO_SPACES_REGION}.digitaloceanspaces.com")
+
+if not DO_SPACES_KEY or not DO_SPACES_SECRET:
+    raise RuntimeError("Missing DigitalOcean Spaces credentials in environment variables")
+
+# Initialize S3 client for DigitalOcean Spaces
+s3_client = boto3.client(
+    's3',
+    region_name=DO_SPACES_REGION,
+    endpoint_url=DO_SPACES_ENDPOINT,
+    aws_access_key_id=DO_SPACES_KEY,
+    aws_secret_access_key=DO_SPACES_SECRET
+)
+
+# ---------------------------------------------------------------------
+# FastAPI app setup
+# ---------------------------------------------------------------------
+app = FastAPI(title="Qwen Image Edit API with Firebase Auth")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ---------------------------------------------------------------------
+# Auth dependency
+# ---------------------------------------------------------------------
+async def verify_firebase_token(request: Request):
+    """Middleware-like dependency to verify Firebase JWT from Authorization header."""
+    auth_header = request.headers.get("Authorization")
+    if not auth_header or not auth_header.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
+
+    id_token = auth_header.split("Bearer ")[1]
+    try:
+        decoded_token = auth.verify_id_token(id_token)
+        request.state.user = decoded_token
+        return decoded_token
+    except Exception as e:
+        raise HTTPException(status_code=401, detail=f"Invalid or expired Firebase token: {e}")
+
+# ---------------------------------------------------------------------
+# Models
+# ---------------------------------------------------------------------
+class HealthResponse(BaseModel):
+    status: str
+    db: str
+    model: str
+
+# --------------------- UTILS ---------------------
+def resize_image_if_needed(img: Image.Image, max_size=(1024, 1024)) -> Image.Image:
+    """
+    Resize image to fit within max_size while keeping aspect ratio.
+    """
+    if img.width > max_size[0] or img.height > max_size[1]:
+        img.thumbnail(max_size, Image.ANTIALIAS)
+    return img
+# ---------------------------------------------------------------------
+# Lazy Gemini Initialization
+# ---------------------------------------------------------------------
+_genai_initialized = False
+def init_gemini():
+    global _genai_initialized, genai
+
+    if _genai_initialized:
+        return
+
+    if not GEMINI_API_KEY:
+        raise RuntimeError("❌ GEMINI_API_KEY not set")
+
+    import google.generativeai as genai
+    genai.configure(api_key=GEMINI_API_KEY)
+
+    _genai_initialized = True
+    
+def expiry(hours: int):
+    """Return UTC datetime for TTL expiration"""
+    return datetime.utcnow() + timedelta(hours=hours)
+    
+def prepare_image(file_bytes: bytes) -> Image.Image:
+    """
+    Open image and resize if larger than 1024x1024
+    """
+    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
+
+    # ✅ MIN SIZE CHECK
+    if img.width < 200 or img.height < 200:
+        raise HTTPException(
+            status_code=400,
+            detail="Image size is below 200x200 pixels. Please upload a larger image."
+        )
+    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
+    img = resize_image_if_needed(img, max_size=(1024, 1024))
+    return img
+
+def upload_to_digitalocean(image_bytes: bytes, folder: str, filename: str) -> str:
+    """
+    Upload image to DigitalOcean Spaces and return the public URL.
+    folder: 'source' or 'results'
+    """
+    key = f"valentine/{folder}/{filename}"
+    try:
+        s3_client.put_object(
+            Bucket=DO_SPACES_BUCKET,
+            Key=key,
+            Body=image_bytes,
+            ContentType="image/jpeg" if filename.endswith('.jpg') else "image/png",
+            ACL='public-read'
+        )
+        url = f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{key}"
+        return url
+    except Exception as e:
+        raise RuntimeError(f"Failed to upload to DigitalOcean Spaces: {e}")
+
+def generate_image_id() -> str:
+    """
+    Generate a random image ID (same format as MongoDB ObjectId for consistency)
+    """
+    return str(uuid.uuid4().hex[:24])
+    
+MAX_COMPRESSED_SIZE = 2 * 1024 * 1024  # 2 MB
+def compress_pil_image_to_2mb(
+    pil_img: Image.Image,
+    max_dim: int = 1280
+) -> bytes:
+    """
+    Resize + compress PIL image to <= 2MB.
+    Returns JPEG bytes.
+    """
+    img = pil_img.convert("RGB")
+
+    # Resize (maintain aspect ratio)
+    img.thumbnail((max_dim, max_dim), Image.LANCZOS)
+
+    quality = 85
+    buffer = io.BytesIO()
+
+    while quality >= 40:
+        buffer.seek(0)
+        buffer.truncate()
+
+        img.save(
+            buffer,
+            format="JPEG",
+            quality=quality,
+            optimize=True,
+            progressive=True
+        )
+
+        if buffer.tell() <= MAX_COMPRESSED_SIZE:
+            break
+
+        quality -= 5
+
+    return buffer.getvalue()
+    
+def run_image_generation(
+    image1: Image.Image,
+    prompt: str,
+    image2: Optional[Image.Image] = None,
+    force_model: Optional[str] = None
+) -> Image.Image:
+    """
+    Unified image generation interface.
+    QWEN  -> merges images if image2 exists
+    GEMINI -> passes images separately
+    """
+
+    effective_model = (force_model or MODEL).upper()
+
+    # ---------------- QWEN ----------------
+    if effective_model == "QWEN":
+
+        # ✅ Merge images ONLY for QWEN
+        if image2:
+            total_width = image1.width + image2.width
+            max_height = max(image1.height, image2.height)
+            merged = Image.new("RGB", (total_width, max_height))
+            merged.paste(image1, (0, 0))
+            merged.paste(image2, (image1.width, 0))
+        else:
+            merged = image1
+
+        return hf_client.image_to_image(
+            image=merged,
+            prompt=prompt,
+            model="Qwen/Qwen-Image-Edit"
+        )
+
+    # ---------------- GEMINI ----------------
+    elif effective_model == "GEMINI":
+        init_gemini()
+        model = genai.GenerativeModel(GEMINI_IMAGE_MODEL)
+
+        parts = [prompt]
+
+        def add_image(img: Image.Image):
+            buf = io.BytesIO()
+            img.save(buf, format="PNG")
+            buf.seek(0)
+            parts.append({
+                "mime_type": "image/png",
+                "data": buf.getvalue()
+            })
+
+        add_image(image1)
+
+        if image2:
+            add_image(image2)
+
+        response = model.generate_content(parts)
+
+        # -------- SAFE IMAGE EXTRACTION --------
+        import base64
+
+        image_bytes = None
+        for candidate in response.candidates:
+            for part in candidate.content.parts:
+                if hasattr(part, "inline_data") and part.inline_data:
+                    data = part.inline_data.data
+                    image_bytes = (
+                        data if isinstance(data, (bytes, bytearray))
+                        else base64.b64decode(data)
+                    )
+                    break
+            if image_bytes:
+                break
+
+        if not image_bytes:
+            raise RuntimeError("Gemini did not return an image")
+
+        img = Image.open(io.BytesIO(image_bytes))
+        img.verify()
+        return Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+    else:
+        raise RuntimeError(f"Unsupported IMAGE_MODEL: {effective_model}")
+
+
+
+# ---------------------------------------------------------------------
+# Endpoints
+# ---------------------------------------------------------------------
+@app.get("/")
+async def root():
+    """Root endpoint"""
+    return {
+        "success": True,
+        "message": "Polaroid,Kiddo,Makeup,Hairstyle API",
+        "data": {
+            "version": "1.0.1",
+            "Product Name":"Beauty Camera - GlowCam AI Studio",
+            "Released By" : "LogicGo Infotech"
+        }
+    }
+    
+@app.get("/health", response_model=HealthResponse)
+def health():
+    """Public health check"""
+    mongo.admin.command("ping")
+    return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
+
+@app.post("/generate")
+async def generate(
+    prompt: str = Form(...),
+    image1: UploadFile = File(...),
+    image2: Optional[UploadFile] = File(None),
+    user_id: Optional[str] = Form(None),
+    category_id: Optional[str] = Form(None),
+    appname: Optional[str] = Form(None),
+    user=Depends(verify_firebase_token)
+):
+    start_time = time.time()
+
+    # -------------------------
+    # 1. VALIDATE & READ IMAGES
+    # -------------------------
+    try:
+        img1_bytes = await image1.read()
+        pil_img1 = prepare_image(img1_bytes)
+        input1_id = generate_image_id()
+        input1_filename = f"{input1_id}_{image1.filename}"
+        input1_url = upload_to_digitalocean(img1_bytes, "source", input1_filename)
+    except Exception as e:
+        raise HTTPException(400, f"Failed to read first image: {e}")
+
+    img2_bytes = None
+    input2_id = None
+    input2_url = None
+    pil_img2 = None
+
+    if image2:
+        try:
+            img2_bytes = await image2.read()
+            pil_img2 = prepare_image(img2_bytes)
+            input2_id = generate_image_id()
+            input2_filename = f"{input2_id}_{image2.filename}"
+            input2_url = upload_to_digitalocean(img2_bytes, "source", input2_filename)
+        except Exception as e:
+            raise HTTPException(400, f"Failed to read second image: {e}")
+
+    # -------------------------
+    # 3. HF INFERENCE
+    # -------------------------
+    try:
+        # --------------------------------------------------
+        # MODEL OVERRIDE BASED ON CATEGORY
+        # --------------------------------------------------
+        force_model = None
+        
+        if category_id == GEMINI_FORCE_CATEGORY_ID:
+            force_model = "GEMINI"
+        
+        pil_output = run_image_generation(
+            image1=pil_img1,
+            image2=pil_img2,
+            prompt=prompt,
+            force_model="GEMINI" if category_id == GEMINI_FORCE_CATEGORY_ID else None
+        )
+
+
+    except Exception as e:
+        response_time_ms = round((time.time() - start_time) * 1000)
+        # logs_collection.insert_one({
+        #     "timestamp": datetime.utcnow(),
+        #     "status": "failure",
+        #     "input1_id": input1_id,
+        #     "input2_id": input2_id if input2_id else None,
+        #     "prompt": prompt,
+        #     "user_email": user.get("email"),
+        #     "error": str(e),
+        #     "response_time_ms": response_time_ms,
+        #     "appname": appname
+        # 
+        if logs_collection is not None:
+            logs_collection.insert_one({
+                "endpoint": "/generate",
+                "status": "fail",
+                "response_time_ms": float(response_time_ms),
+                "timestamp": datetime.utcnow(),
+                "appname": appname if appname else "None",
+                "error": str(e)
+            })
+        raise HTTPException(500, f"Inference failed: {e}")
+
+    # -------------------------
+    # 5. SAVE OUTPUT IMAGE
+    # -------------------------
+    output_object_id = ObjectId()
+    output_image_id = str(output_object_id)
+    out_buf = io.BytesIO()
+    pil_output.save(out_buf, format="PNG")
+    out_bytes = out_buf.getvalue()
+    
+    out_filename = f"{output_image_id}_result.png"
+    out_url = upload_to_digitalocean(out_bytes, "results", out_filename)
+    try:
+        fs.put(
+            out_bytes,
+            _id=output_object_id,
+            filename=out_filename,
+            content_type="image/png",
+            metadata={
+                "user_email": user.get("email"),
+                "prompt": prompt,
+                "created_at": datetime.utcnow()
+            }
+        )
+    except Exception as e:
+        raise HTTPException(500, f"Failed to store output image in GridFS: {e}")
+
+    # -------------------------
+    # 5b. SAVE COMPRESSED IMAGE
+    # -------------------------
+    compressed_bytes = compress_pil_image_to_2mb(
+        pil_output,
+        max_dim=1280
+    )
+    
+    compressed_filename = f"{output_image_id}_compressed.jpg"
+    compressed_url = upload_to_digitalocean(compressed_bytes, "results", compressed_filename)
+
+
+    response_time_ms = round((time.time() - start_time) * 1000)
+
+    # -------------------------
+    # 6. LOG SUCCESS
+    # -------------------------
+    # logs_collection.insert_one({
+    #     "timestamp": datetime.utcnow(),
+    #     "status": "success",
+    #     "input1_id": input1_id,
+    #     "input2_id": input2_id if input2_id else None,
+    #     "output_id": output_image_id,
+    #     "prompt": prompt,
+    #     "user_email": user.get("email"),
+    #     "response_time_ms": response_time_ms,
+    #     "appname": appname
+    # })
+    if logs_collection is not None:
+        logs_collection.insert_one({
+            "endpoint": "/generate",
+            "status": "success",
+            "response_time_ms": float(response_time_ms),
+            "timestamp": datetime.utcnow(),
+            "appname": appname if appname else "None",
+            "error": None
+        })
+    return JSONResponse({
+        "output_image_id": output_image_id,
+        "user": user.get("email"),
+        "response_time_ms": response_time_ms,
+        "Compressed_Image_URL": compressed_url
+    })
+
+
+# Image endpoint removed - images are now stored directly in DigitalOcean Spaces
+# and are publicly accessible via the Compressed_Image_URL returned in the /generate response
+@app.get("/image/{image_id}")
+async def download_image(
+    image_id: str
+):
+    try:
+        file_obj_id = ObjectId(image_id)
+    except Exception:
+        raise HTTPException(status_code=400, detail="Invalid image id format")
+
+    def _read_from_gridfs():
+        grid_out = fs.get(file_obj_id)
+        return grid_out.read(), grid_out.filename, getattr(grid_out, "content_type", "image/png")
+
+    try:
+        file_bytes, filename, content_type = await run_in_threadpool(_read_from_gridfs)
+    except NoFile:
+        raise HTTPException(status_code=404, detail="Image not found")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to download image: {e}")
+
+    return StreamingResponse(
+        io.BytesIO(file_bytes),
+        media_type=content_type or "image/png",
+        headers={"Content-Disposition": f'attachment; filename="{filename or f"{image_id}.png"}"'}
+    )
+
+# ---------------------------------------------------------------------
+# Run locally
+# ---------------------------------------------------------------------
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)