from fastapi import FastAPI, HTTPException, Header
from huggingface_hub import InferenceClient
from fastapi.responses import Response
import io
import os

app = FastAPI()

HF_TOKEN = os.environ.get("HF_TOKEN")
client = InferenceClient(token=HF_TOKEN)

API_ACCESS_SECRET = "h7-vision-ml-secure-v1"

@app.get("/")
def read_root():
    return {"status": "Vision AI Engine is running", "model": "FLUX.1-schnell"}

@app.post("/generate")
async def generate_image(prompt: str, x_api_key: str = Header(None)):
    if x_api_key != API_ACCESS_SECRET:
        raise HTTPException(status_code=403, detail="Unauthorized Access to h7 Engine")

    try:
        image = client.text_to_image(prompt, model="black-forest-labs/FLUX.1-schnell")
        
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        
        return Response(content=img_byte_arr.getvalue(), media_type="image/png")
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))