app.py

import gradio as gr
import requests
import json
import time
from better_profanity import profanity

url = "https://nexra.aryahcr.cc/api/image/complements"
headers = {
    "Content-Type": "application/json"
}

profanity.load_censor_words()
no_nsfw_image_path = "no_nsfw.png"  

def is_nsfw(text):
    return profanity.contains_profanity(text)

def generate_image(prompt, model):
    if is_nsfw(prompt):
        return no_nsfw_image_path  

    data = {
        "prompt": prompt,
        "model": model,
    }
    try:
        response = requests.post(url, headers=headers, data=json.dumps(data))
        
        if response.status_code == 200:
            response_data = response.json()
            image_id = response_data.get("id")
            
            if not image_id:
                return no_nsfw_image_path 

            while True:
                status_response = requests.get(f"{url}/{image_id}")
                
                if status_response.status_code == 200:
                    status_data = status_response.json()
                    status = status_data.get("status")
                    
                    if status == "completed":
                        images = status_data.get("images")
                        if images and isinstance(images, list):
                            image_url = images[0]
                            return image_url
                        else:
                            return no_nsfw_image_path  
                    elif status == "error":
                        return no_nsfw_image_path
                    elif status == "not_found":
                        return no_nsfw_image_path
                    elif status == "pending":
                        time.sleep(1)
                    else:
                        return no_nsfw_image_path
                else:
                    return no_nsfw_image_path
        else:
            return no_nsfw_image_path
    except json.JSONDecodeError:
        return no_nsfw_image_path
    except Exception as e:
        return no_nsfw_image_path

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Enter prompt"),
        gr.Radio(["dalle2"], label="Select Model")
    ],
    outputs=gr.Image(type="filepath"),
    title="DALLE2 Generation",
    description="DALLE-2 Generation from Nextra API, NOTE: DALL-E 3 is not supported yet."
)
iface.launch()