app.py

from fastapi import FastAPI, Response
from pydantic import BaseModel
import requests

API_URL = "https://ai.sitius.ir/v1/generate/"

app = FastAPI()

class GenerationRequest(BaseModel):
    prompt: str
    model: str
    negative_prompt: str
    steps: int
    cfg_scale: int
    sampler: str
def query(payload):
    response = requests.post(API_URL, json=payload)
    response.raise_for_status()
    return response.content

@app.post("/generate")
def generate_image(request: GenerationRequest):
    try:
        image_bytes = query({
            "prompt": request.prompt,
            "model": request.model, 
            "negative_prompt": request.negative,
            "steps": 40,
            "cfg_scale": 7,
            "sampler": "sampler",

        })
        return Response(content=image_bytes, media_type="image/png")
    
    except Exception as e:
        return {"error": str(e)}