from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from gradio_client import Client
from PIL import Image
import io
import os

app = FastAPI()

# Allow CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Connect to the model
client = Client("NihalGazi/FLUX-Pro-Unlimited")

@app.get("/")
def root():
    return {"message": "Welcome to the JanitorAI Generator API! Use /gen?prompt={prompt} to request an image"}

@app.get("/gen")
def generate_image(
    prompt: str = Query(..., description="Prompt for image generation"),
    basemodel: str = "black-forest-labs/FLUX.1-DEV",
    width: int = 1280,
    height: int = 768,
    scales: int = 8,
    steps: int = 8,
    seed: int = -1,
    upscale_factor: str = "2",
    process_upscale: bool = False,
    lora_model: str = "enhanceaiteam/Flux-uncensored",
    process_lora: bool = False
):
    try:
        result = client.predict(
    		prompt=prompt,
    		width=1280,
    		height=1280,
    		seed=0,
    		randomize=True,
    		server_choice="NSFW-Core: Uncensored Server 2",
    		api_name="/generate_image"
        )
        print(result)

        # result is a list of local file paths
        if isinstance(result, list) and result:
            image_path = result[0]
        else:
            raise ValueError("No image returned from model.")

        # Read the image from disk
        with open(image_path, "rb") as img_file:
            img_bytes = img_file.read()

        return StreamingResponse(io.BytesIO(img_bytes), media_type="image/png")

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))