app.py

import os
import requests
import json
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 未在环境变量中设置")

# 安全分类配置
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"安全模型加载失败: {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):
    """将图片转换为Base64格式"""
    with open(file_path, "rb") as f:
        file_ext = Path(file_path).suffix.lower()[1:]
        mime_type = f"image/{file_ext}" if file_ext in ["jpeg", "jpg", "png"
                                                        ] else "image/jpeg"
        return f"data:{mime_type};base64,{base64.b64encode(f.read()).decode()}"


# 核心生成逻辑
def generate_image(image_file, prompt, seed, session_id, enable_safety=True):
    try:
        # 安全检查
        if enable_safety:
            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 "❌ 请求过于频繁，请稍后再试", 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("请上传图片文件")
        elif not Path(image_file).exists():
            error_messages.append("文件不存在")
        elif Path(image_file).suffix.lower()[1:] not in ["jpg", "jpeg", "png"]:
            error_messages.append("仅支持JPG/PNG格式")
        if not prompt.strip():
            error_messages.append("提示语不能为空")
        if error_messages:
            error_img = create_error_image(" | ".join(error_messages))
            yield "❌ 输入验证失败", error_img, ""
            return

        # 转换为Base64
        try:
            base64_image = image_to_base64(image_file)
        except Exception as e:
            error_img = create_error_image(f"文件处理失败: {str(e)}")
            yield "❌ 文件处理失败", error_img, ""
            return

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

        # 提交请求
        response = requests.post(
            "https://api.wavespeed.ai/api/v2/wavespeed-ai/hidream-e1-full",
            headers=headers,
            json=payload,
            timeout=30)
        response.raise_for_status()

        # 处理响应
        request_id = response.json()["data"]["id"]
        result_url = f"https://api.wavespeed.ai/api/v2/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"🎉 生成成功! 耗时 {elapsed:.1f}s", image_url, image_url
                return
            elif status == "failed":
                raise Exception(data.get("error", "Unknown error"))
            else:
                yield f"⏳ 当前状态: {status.capitalize()}...", None, None

        raise Exception("生成超时")

    except Exception as e:
        error_img = create_error_image(str(e))
        yield f"❌ 生成失败: {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("# 🖼️Hidream-E1-Full Live On  Wavespeed Ai")
    gr.Markdown("HiDream-E1 is an image editing model built on HiDream-I1.")

    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 an English 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("开始生成", 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])

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