app.py

import os
import requests
import time
import threading
import uuid
import base64
from pathlib import Path
from dotenv import load_dotenv
import gradio as gr
import random
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoTokenizer, AutoModelForSequenceClassification

load_dotenv()
API_KEY = os.getenv("WAVESPEED_API_KEY")
if not API_KEY:
    raise ValueError("WAVESPEED_API_KEY is not set in environment variables")

MODEL_URL = "TostAI/nsfw-text-detection-large"
CLASS_NAMES = {0: "✅ SAFE", 1: "⚠️ QUESTIONABLE", 2: "🚫 UNSAFE"}

try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL)
except Exception as e:
    raise RuntimeError(f"Failed to load safety model: {str(e)}")


class SessionManager:
    _instances = {}
    _lock = threading.Lock()

    @classmethod
    def get_session(cls, session_id):
        with cls._lock:
            if session_id not in cls._instances:
                cls._instances[session_id] = {
                    'count': 0,
                    'history': [],
                    'last_active': time.time()
                }
            return cls._instances[session_id]

    @classmethod
    def cleanup_sessions(cls):
        with cls._lock:
            now = time.time()
            expired = [
                k for k, v in cls._instances.items()
                if now - v['last_active'] > 3600
            ]
            for k in expired:
                del cls._instances[k]


class RateLimiter:

    def __init__(self):
        self.clients = {}
        self.lock = threading.Lock()

    def check(self, client_id):
        with self.lock:
            now = time.time()
            if client_id not in self.clients:
                self.clients[client_id] = {'count': 1, 'reset': now + 3600}
                return True
            if now > self.clients[client_id]['reset']:
                self.clients[client_id] = {'count': 1, 'reset': now + 3600}
                return True
            if self.clients[client_id]['count'] >= 20:
                return False
            self.clients[client_id]['count'] += 1
            return True


session_manager = SessionManager()
rate_limiter = RateLimiter()


def create_error_image(message):
    img = Image.new("RGB", (512, 512), "#ffdddd")
    try:
        font = ImageFont.truetype("arial.ttf", 24)
    except:
        font = ImageFont.load_default()
    draw = ImageDraw.Draw(img)
    text = f"Error: {message[:60]}..." if len(message) > 60 else message
    draw.text((50, 200), text, fill="#ff0000", font=font)
    return img


@torch.no_grad()
def classify_prompt(prompt):
    inputs = tokenizer(prompt,
                       return_tensors="pt",
                       truncation=True,
                       max_length=512)
    outputs = model(**inputs)
    return torch.argmax(outputs.logits).item()


def image_to_base64(file_path):
    with open(file_path, "rb") as f:
        return base64.b64encode(f.read()).decode()


def generate_image(image_file,
                   prompt,
                   seed,
                   session_id,
                   enable_safety_checker=True):
    try:
        if enable_safety_checker:
            safety_level = classify_prompt(prompt)
            if safety_level != 0:
                error_img = create_error_image(CLASS_NAMES[safety_level])
                yield f"❌ Blocked: {CLASS_NAMES[safety_level]}", error_img, ""
                return

        if not rate_limiter.check(session_id):
            error_img = create_error_image(
                "Hourly limit exceeded (20 requests)")
            yield "❌ Too many requests, please try again later", error_img, ""
            return

        session = session_manager.get_session(session_id)
        session['last_active'] = time.time()
        session['count'] += 1

        error_messages = []
        if not image_file:
            error_messages.append("Please upload an image file")
        elif not Path(image_file).exists():
            error_messages.append("File does not exist")
        if not prompt.strip():
            error_messages.append("Prompt cannot be empty")
        if error_messages:
            error_img = create_error_image(" | ".join(error_messages))
            yield "❌ Input validation failed", error_img, ""
            return

        try:
            base64_image = image_to_base64(image_file)
        except Exception as e:
            error_img = create_error_image(f"File processing failed: {str(e)}")
            yield "❌ File processing failed", error_img, ""
            return

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}",
        }
        payload = {
            "enable_safety_checker": enable_safety_checker,
            "image": base64_image,
            "prompt": prompt,
            "seed": int(seed) if seed != -1 else random.randint(0, 999999)
        }

        response = requests.post(
            "https://api.wavespeed.ai/api/v3/wavespeed-ai/flux-kontext-dev-ultra-fast",
            headers=headers,
            json=payload,
            timeout=30)
        response.raise_for_status()

        request_id = response.json()["data"]["id"]
        result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
        start_time = time.time()

        for _ in range(60):
            time.sleep(1)
            resp = requests.get(result_url, headers=headers)
            resp.raise_for_status()

            data = resp.json()["data"]
            status = data["status"]

            if status == "completed":
                elapsed = time.time() - start_time
                image_url = data["outputs"][0]
                session["history"].append(image_url)
                yield f"🎉 Generation successful! Time taken {elapsed:.1f}s", image_url, image_url
                return
            elif status == "failed":
                raise Exception(data.get("error", "Unknown error"))
            else:
                yield f"⏳ Current status: {status.capitalize()}...", None, None

        raise Exception("Generation timed out")

    except Exception as e:
        error_img = create_error_image(str(e))
        yield f"❌ Generation failed: {str(e)}", error_img, ""


def cleanup_task():
    while True:
        session_manager.cleanup_sessions()
        time.sleep(3600)


with gr.Blocks(theme=gr.themes.Soft(),
               css="""
    .status-box { padding: 10px; border-radius: 5px; margin: 5px; }
    .safe { background: #e8f5e9; border: 1px solid #a5d6a7; }
    .warning { background: #fff3e0; border: 1px solid #ffcc80; }
    .error { background: #ffebee; border: 1px solid #ef9a9a; }
    """) as app:

    session_id = gr.State(str(uuid.uuid4()))

    gr.Markdown("# 🖼️FLUX Kontext Dev Ultra Fast Live On Wavespeed AI")
    gr.Markdown(
        "FLUX Kontext Dev is a new SOTA image editing model published by Black Forest Labs. We have deployed it on [WaveSpeedAI](https://wavespeed.ai/) for ultra-fast image editing. You can use it to edit images in various styles, add objects, or even change the mood of the image. It supports both text prompts and image inputs."
    )
    gr.Markdown(
        "[FLUX Kontext Dev on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev)"
    )
    gr.Markdown(
        "[FLUX Kontext Dev Ultra Fast on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-ultra-fast)"
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_file = gr.Image(label="Upload Image",
                                  type="filepath",
                                  sources=["upload"],
                                  interactive=True,
                                  image_mode="RGB")
            prompt = gr.Textbox(label="Prompt",
                                placeholder="Please enter your prompt...",
                                lines=3)
            seed = gr.Number(label="seed",
                             value=-1,
                             minimum=-1,
                             maximum=999999,
                             step=1)
            random_btn = gr.Button("random🎲seed", variant="secondary")
            enable_safety = gr.Checkbox(label="🔒 Enable Safety Checker",
                                        value=True,
                                        interactive=False)
        with gr.Column(scale=1):
            output_image = gr.Image(label="Generated Result")
            output_url = gr.Textbox(label="Image URL",
                                    interactive=True,
                                    visible=False)
            status = gr.Textbox(label="Status", elem_classes=["status-box"])
            submit_btn = gr.Button("Start Generation", variant="primary")
    gr.Examples(
        examples=[
            [
                "Convert the image into Claymation style.",
                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
            ],
            [
                "Convert the image into a Ghibli style.",
                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flux_ip_adapter_input.jpg"
            ],
            [
                "Add sunglasses to the face of the girl.",
                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_mask_girl2.png"
            ],
            # [
            #     'Convert the image into an ink sketch style.',
            #     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
            # ],
            # [
            #     'Add a butterfly to the scene.',
            #     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_depth_result.png"
            # ]
        ],
        inputs=[prompt, image_file],
        label="Examples")

    random_btn.click(fn=lambda: random.randint(0, 999999), outputs=seed)

    submit_btn.click(
        generate_image,
        inputs=[image_file, prompt, seed, session_id, enable_safety],
        outputs=[status, output_image, output_url],
        api_name=False,
    )

if __name__ == "__main__":
    threading.Thread(target=cleanup_task, daemon=True).start()
    app.queue(max_size=8).launch(
        server_name="0.0.0.0",
        share=False,
    )