Viewing File: /home/ubuntu/combine_ai/text_to_image/main_image.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, field_validator
from openai import OpenAI
import os
import openai
from dotenv import load_dotenv
from datetime import datetime, timedelta
from typing import Optional
import uvicorn
import re

app = FastAPI()

load_dotenv()

client = OpenAI()
openai.api_key = os.getenv("OPENAI_API_KEY")

last_request_time = {}


class ImageRequest(BaseModel):
    user_id: str
    model: str
    individual_prompt: str
    predefined_prompt: Optional[str] = None
    lighting_style: Optional[str] = None
    size: str

    @field_validator('individual_prompt')
    def validate_prompt(cls, v):
        if not re.search("[a-zA-Z]", v):
            raise ValueError("Prompt must contain meaningful words and not just numbers or special characters.")
        return v

def check_cooldown(user_id: str):
    now = datetime.now()
    if user_id in last_request_time:
        if now - last_request_time[user_id] < timedelta(minutes=2):
            return False
    last_request_time[user_id] = now
    return True


@app.post("/generate_image/")
async def generate_image(request: ImageRequest):
    if not check_cooldown(request.user_id):
        raise HTTPException(status_code=429, detail="please wait for 2 minutes before making another request")

    parts = [
        part for part in [
            request.model,
            request.predefined_prompt,
            request.lighting_style,
            request.individual_prompt
        ] if part is not None
    ]
    combined_prompt = ", ".join(parts)
    print(combined_prompt)
    try:
        response = client.images.generate(
            model="dall-e-3",
            prompt=combined_prompt,
            size="1024x1024",
            quality="standard",
            n=1,
        )
        image_url = response.data[0].url
        return {"url": image_url, "success": True}
    except Exception as e:
        return {"success": False, "error": str(e)}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8002, ssl_keyfile="privkey.pem", ssl_certfile="fullchain.pem")