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
import io
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 decode_base64_to_image(base64_str):
    image_data = base64.b64decode(base64_str)
    return Image.open(io.BytesIO(image_data))


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, "", None
                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, "", None
            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, "", None
            return

        try:
            base64_image = image_to_base64(image_file)
            input_image = decode_base64_to_image(base64_image)
        except Exception as e:
            error_img = create_error_image(f"File processing failed: {str(e)}")
            yield "❌ File processing failed", error_img, "", None
            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
                output_url = data["outputs"][0]
                session["history"].append(output_url)
                yield f"🎉 Generation successful! Time taken {elapsed:.1f}s", output_url, output_url, update_recent_gallery(prompt, input_image, output_url)
                return
            elif status == "failed":
                raise Exception(data.get("error", "Unknown error"))
            else:
                yield f"⏳ Current status: {status.capitalize()}...", None, 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, "", None


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


# Store recent generations
recent_generations = []

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")
    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)"
        "- [FLUX Kontext Dev LoRA on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-lora)"
        "- [FLUX Kontext Dev Ultra Fast on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-ultra-fast)"
        "- [FLUX Kontext Dev LoRA Ultra Fast on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-lora-ultra-fast)"
    )
    gr.Markdown(
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="Prompt",
                                placeholder="Please enter your prompt...",
                                lines=3)
            image_file = gr.Image(label="Upload Image",
                                  type="filepath",
                                  sources=["upload"],
                                  interactive=True,
                                  image_mode="RGB")
            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):
            status = gr.Textbox(label="Status", elem_classes=["status-box"])
            output_image = gr.Image(label="Generated Result")
            output_url = gr.Textbox(label="Image URL",
                                    interactive=True,
                                    visible=False)
            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 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")

    with gr.Accordion("Recent Generations (last 16)", open=False):
        recent_gallery = gr.Gallery(label="Prompt and Output",
                                    columns=3,
                                    interactive=False)

    def get_recent_gallery_items():
        gallery_items = []
        for r in reversed(recent_generations):
            if any(x is None for x in r.values()):
                continue
            gallery_items.append((r["input"], f"Input: {r['prompt']}"))
            gallery_items.append((r["output"], f"Output: {r['prompt']}"))
        return gr.update(value=gallery_items)

    def update_recent_gallery(prompt, input_image, output_image):
        recent_generations.append({
            "prompt": prompt,
            "input": input_image,
            "output": output_image,
        })
        if len(recent_generations) > 16:
            recent_generations.pop(0)
        return get_recent_gallery_items()

    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, recent_gallery],
        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,
    )