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