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

from fastapi import FastAPI, HTTPException, Request
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 re
import uvicorn

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)
    dummy_url = "https://combineai-s3.s3.amazonaws.com/DALL%C2%B7E%202024-06-20%2013.44.12%20-%20A%20serene%20landscape%20depicting%20a%20tranquil%20lake%20with%20crystal%20clear%20water%20reflecting%20the%20sky%2C%20surrounded%20by%20lush%20green%20forests%20under%20a%20bright%20blue%20sky%20wit.jpg"
    '''
    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)}
    '''
    return {"url": dummy_url, "success": True}


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