|
|
import os |
|
|
import requests |
|
|
import json |
|
|
import time |
|
|
import random |
|
|
import base64 |
|
|
import uuid |
|
|
import threading |
|
|
from pathlib import Path |
|
|
from dotenv import load_dotenv |
|
|
import gradio as gr |
|
|
import torch |
|
|
from PIL import Image, ImageDraw, ImageFont |
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
|
|
|
|
load_dotenv() |
|
|
|
|
|
MODEL_URL = "TostAI/nsfw-text-detection-large" |
|
|
CLASS_NAMES = {0: "✅ SAFE", 1: "⚠️ QUESTIONABLE", 2: "🚫 UNSAFE"} |
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_URL) |
|
|
model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL) |
|
|
|
|
|
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", (832, 480), "#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) |
|
|
img.save("error.jpg") |
|
|
return "error.jpg" |
|
|
|
|
|
def classify_prompt(prompt): |
|
|
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512) |
|
|
with torch.no_grad(): |
|
|
outputs = model(**inputs) |
|
|
return torch.argmax(outputs.logits).item() |
|
|
|
|
|
def image_to_base64(file_path): |
|
|
try: |
|
|
with open(file_path, "rb") as image_file: |
|
|
ext = Path(file_path).suffix.lower().lstrip('.') |
|
|
mime_map = { |
|
|
'jpg': 'jpeg', |
|
|
'jpeg': 'jpeg', |
|
|
'png': 'png', |
|
|
'webp': 'webp', |
|
|
'gif': 'gif' |
|
|
} |
|
|
mime_type = mime_map.get(ext, 'jpeg') |
|
|
|
|
|
raw_data = image_file.read() |
|
|
encoded = base64.b64encode(raw_data) |
|
|
missing_padding = len(encoded) % 4 |
|
|
if missing_padding: |
|
|
encoded += b'=' * (4 - missing_padding) |
|
|
|
|
|
return f"data:image/{mime_type};base64,{encoded.decode('utf-8')}" |
|
|
except Exception as e: |
|
|
raise ValueError(f"Base64编码失败: {str(e)}") |
|
|
|
|
|
def generate_video( |
|
|
image, |
|
|
prompt, |
|
|
enable_safety, |
|
|
flow_shift, |
|
|
guidance_scale, |
|
|
negative_prompt, |
|
|
seed, |
|
|
size, |
|
|
session_id |
|
|
): |
|
|
|
|
|
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("每小时限制20次请求") |
|
|
yield "❌ 请求过于频繁，请稍后再试", error_img |
|
|
return |
|
|
|
|
|
session = session_manager.get_session(session_id) |
|
|
session['last_active'] = time.time() |
|
|
session['count'] += 1 |
|
|
|
|
|
API_KEY = os.getenv("WAVESPEED_API_KEY") |
|
|
if not API_KEY: |
|
|
error_img = create_error_image("API密钥缺失") |
|
|
yield "❌ Error: Missing API Key", error_img |
|
|
return |
|
|
|
|
|
try: |
|
|
base64_image = image_to_base64(image) |
|
|
except Exception as e: |
|
|
error_img = create_error_image(str(e)) |
|
|
yield f"❌ 文件上传失败: {str(e)}", error_img |
|
|
return |
|
|
|
|
|
payload = { |
|
|
"enable_safety_checker": enable_safety, |
|
|
"flow_shift": flow_shift, |
|
|
"guidance_scale": guidance_scale, |
|
|
"image": base64_image, |
|
|
"negative_prompt": negative_prompt, |
|
|
"prompt": prompt, |
|
|
"seed": seed if seed != -1 else random.randint(0, 999999), |
|
|
"size": size |
|
|
} |
|
|
|
|
|
headers = { |
|
|
"Content-Type": "application/json", |
|
|
"Authorization": f"Bearer {API_KEY}", |
|
|
} |
|
|
|
|
|
try: |
|
|
response = requests.post( |
|
|
"https://api.wavespeed.ai/api/v2/wavespeed-ai/hunyuan-custom-ref2v-480p", |
|
|
headers=headers, |
|
|
data=json.dumps(payload) |
|
|
) |
|
|
|
|
|
if response.status_code != 200: |
|
|
error_img = create_error_image(response.text) |
|
|
yield f"❌ API错误 ({response.status_code}): {response.text}", error_img |
|
|
return |
|
|
|
|
|
request_id = response.json()["data"]["id"] |
|
|
yield f"✅ 任务已提交 (ID: {request_id})", None |
|
|
except Exception as e: |
|
|
error_img = create_error_image(str(e)) |
|
|
yield f"❌ 连接错误: {str(e)}", error_img |
|
|
return |
|
|
|
|
|
result_url = f"https://api.wavespeed.ai/api/v2/predictions/{request_id}/result" |
|
|
start_time = time.time() |
|
|
|
|
|
while True: |
|
|
time.sleep(0.5) |
|
|
try: |
|
|
response = requests.get(result_url, headers=headers) |
|
|
if response.status_code != 200: |
|
|
error_img = create_error_image(response.text) |
|
|
yield f"❌ 轮询错误 ({response.status_code}): {response.text}", error_img |
|
|
return |
|
|
|
|
|
data = response.json()["data"] |
|
|
status = data["status"] |
|
|
|
|
|
if status == "completed": |
|
|
elapsed = time.time() - start_time |
|
|
video_url = data['outputs'][0] |
|
|
session["history"].append(video_url) |
|
|
yield (f"🎉 完成! 耗时 {elapsed:.1f}秒\n" |
|
|
f"下载链接: {video_url}"), video_url |
|
|
return |
|
|
|
|
|
elif status == "failed": |
|
|
error_img = create_error_image(data.get('error', '未知错误')) |
|
|
yield f"❌ 任务失败: {data.get('error', '未知错误')}", error_img |
|
|
return |
|
|
|
|
|
else: |
|
|
yield f"⏳ 状态: {status.capitalize()}...", None |
|
|
|
|
|
except Exception as e: |
|
|
error_img = create_error_image(str(e)) |
|
|
yield f"❌ 轮询失败: {str(e)}", error_img |
|
|
return |
|
|
|
|
|
def cleanup_task(): |
|
|
while True: |
|
|
session_manager.cleanup_sessions() |
|
|
time.sleep(3600) |
|
|
|
|
|
with gr.Blocks( |
|
|
theme=gr.themes.Soft(), |
|
|
css=""" |
|
|
.video-preview { max-width: 600px !important; } |
|
|
.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("# 🌊Hunyuan-Custom-Ref2v Run On [WaveSpeedAI](https://wavespeed.ai/)") |
|
|
gr.Markdown("""HunyuanCustom, a multi-modal, conditional, and controllable generation model centered on subject consistency, built upon the Hunyuan Video generation framework. It enables the generation of subject-consistent videos conditioned on text, images, audio, and video inputs.""") |
|
|
|
|
|
with gr.Row(): |
|
|
with gr.Column(scale=1): |
|
|
img_input = gr.Image(type="filepath", label="Input Image") |
|
|
prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Prompt...") |
|
|
negative_prompt = gr.Textbox(label="Negative Prompt", lines=2) |
|
|
size = gr.Dropdown(["832*480", "480*832"], value="832*480", label="Size") |
|
|
seed = gr.Number(-1, label="Seed") |
|
|
random_seed_btn = gr.Button("Random🎲Seed", variant="secondary") |
|
|
guidance = gr.Slider(1, 20, value=7.5, step=0.1, label="Guidance") |
|
|
flow_shift = gr.Slider(1, 20, value=13, step=1, label="Shift") |
|
|
enable_safety = gr.Checkbox(True, label="Enable Safety Checker", interactive=False) |
|
|
|
|
|
with gr.Column(scale=1): |
|
|
video_output = gr.Video(label="Video Output", format="mp4", interactive=False, elem_classes=["video-preview"]) |
|
|
generate_btn = gr.Button("Generate", variant="primary") |
|
|
status_output = gr.Textbox(label="status", interactive=False, lines=4) |
|
|
|
|
|
gr.Examples( |
|
|
examples=[ |
|
|
[ |
|
|
"A dog is chasing a cat in the park. ", |
|
|
"https://github.com/Tencent/HunyuanCustom/blob/main/assets/images/seg_poodle.png?raw=true" |
|
|
], |
|
|
[ |
|
|
"A single person, in the dressing room. A woman is holding a lipstick, trying it on, and introducing it. ", |
|
|
"https://github.com/Tencent/HunyuanCustom/blob/main/assets/images/seg_boy.png?raw=true" |
|
|
], |
|
|
[ |
|
|
"A man is drinking Moutai in the pavilion. ", |
|
|
"https://github.com/Tencent/HunyuanCustom/blob/main/assets/images/seg_man_03.png?raw=true" |
|
|
], |
|
|
[ |
|
|
"A woman is boxing with a panda, and they are at a stalemate. ", |
|
|
"https://github.com/Tencent/HunyuanCustom/blob/main/assets/images/seg_woman_01.png?raw=true" |
|
|
] |
|
|
], |
|
|
inputs=[prompt, img_input], |
|
|
label="Examples Prompt", |
|
|
examples_per_page=3 |
|
|
) |
|
|
|
|
|
random_seed_btn.click( |
|
|
fn=lambda: random.randint(0, 999999), |
|
|
outputs=seed |
|
|
) |
|
|
|
|
|
generate_btn.click( |
|
|
generate_video, |
|
|
inputs=[ |
|
|
img_input, |
|
|
prompt, |
|
|
enable_safety, |
|
|
flow_shift, |
|
|
guidance, |
|
|
negative_prompt, |
|
|
seed, |
|
|
size, |
|
|
session_id |
|
|
], |
|
|
outputs=[status_output, video_output] |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
threading.Thread(target=cleanup_task, daemon=True).start() |
|
|
app.queue(max_size=4).launch( |
|
|
server_name="0.0.0.0", |
|
|
max_threads=16, |
|
|
share=False |
|
|
) |
