from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
import requests
from urllib.parse import quote
import asyncio
from io import BytesIO
import os

url = os.environ.get("url")

app = FastAPI()

def create_job(prompt, model, sampler, seed, neg):
    if model is None:
        model = 'Realistic_Vision_V5.0.safetensors [614d1063]'
    if sampler is None:
        sampler = 'DPM++ 2M Karras'
    if seed is None:
        seed = '-1'
    if neg is None:
        neg = "(long list of negative prompts removed for brevity)"
    url = f'https://api.{url}.com/generate'
    params = {
        'new': 'true',
        'prompt': quote(prompt),
        'model': model,
        'negative_prompt': neg,
        'steps': '100',
        'cfg': '9.5',
        'seed': seed,
        'sampler': sampler,
        'upscale': 'True',
        'aspect_ratio': 'square'
    }
    response = requests.get(url, params=params)
    response.raise_for_status()
    return response.json()['job']

@app.get("/generate_image")
async def generate_image(prompt: str, model: str = None, sampler: str = None, seed: str = None, neg: str = None):
    job_id = create_job(prompt, model, sampler, seed, neg)
    url = f'https://api.{url}.com/job/{job_id}'
    headers = {'accept': '*/*'}
    
    while True:
        response = requests.get(url=url, headers=headers)
        response.raise_for_status()
        job_response = response.json()
        if job_response['status'] == 'succeeded':
            image_url = f'https://images.{url}.xyz/{job_id}.png'
            image_response = requests.get(image_url)
            image_response.raise_for_status()
            image = BytesIO(image_response.content)
            return StreamingResponse(image, media_type='image/png')
        await asyncio.sleep(2)  # Add a delay to prevent excessive requests

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000, debug=True)