app.py

import gradio as gr 
import numpy as np
import random
import torch
from PIL import Image
import time
from datetime import datetime
from optimum.intel import OVStableDiffusionPipeline
import re
from functools import lru_cache

# Đảm bảo tính tái tạo
torch.manual_seed(0)
np.random.seed(0)
random.seed(0)

# Tối ưu hóa bộ nhớ và tốc độ
torch.set_num_threads(1)  # Giảm xung đột thread
torch.backends.mkldnn.enabled = True  # Kích hoạt MKL-DNN
torch.backends.openmp.enabled = True

print("🔧 Loading OpenVINO pipeline: HARRY07979/sd-v1-5-openvino")

try:
    pipeline = OVStableDiffusionPipeline.from_pretrained(
        "HARRY07979/stable-diffusion-v1-5-openvino",
        safety_checker=None,
        feature_extractor=None,
        torch_dtype=torch.float32
    )
except Exception as e:
    print(f"Error loading pipeline from remote repository: {e}")
    try:
        pipeline = OVStableDiffusionPipeline.from_pretrained(
            "/path/to/local/model",
            safety_checker=None,
            feature_extractor=True,
            torch_dtype=torch.float32,
            local_files_only=True,
        )
    except Exception as e:
        print(f"Error loading local model: {e}")
        raise RuntimeError("Failed to load the pipeline.")

# Tối ưu hóa pipeline
pipeline.to("cpu")
# Loại bỏ dòng pipeline.model.half() vì không hỗ trợ với OVStableDiffusionPipeline
pipeline.compile()  # Biên dịch mô hình nếu hỗ trợ

# Tối ưu hóa từ điển NSFW - sử dụng set để tìm kiếm nhanh hơn
NSFW_HIGH = {
    "nude", "naked", "sex", "porn", "xxx", "fuck", "dick", "cock", "pussy", "vagina", "penis",
    "boobs", "tits", "breasts", "bra", "panties", "underwear", "lingerie", "orgasm", "cum",
    "blowjob", "handjob", "masturbate", "rape", "gangbang", "incest", "hentai", "lewd",
    "erotic", "kinky", "bondage", "bdsm", "squirt", "creampie", "threesome", "orgy", "yaoi",
    "yuri", "futanari", "cunnilingus", "fellatio", "anal", "paizuri", "bukkake", "guro",
    "vore", "tentacle", "netorare", "cuckold", "exhibitionism", "voyeurism", "poop", "pee",
    "poo", "shit", "piss", "scat", "diarrhea", "vomit", "gore", "blood", "murder",
    "torture", "suicide", "decapitation", "mutilation", "drugs", "cocaine", "heroin",
    "lsd", "ecstasy", "vlxx"
}

NSFW_MEDIUM = {
    "bikini", "swimwear", "sexy", "succubus", "leather", "latex", "stockings", "miniskirt",
    "cleavage", "thighs", "ass", "butt", "skirt", "dress", "topless", "wet", "moaning",
    "spread", "legs apart", "tight", "revealing", "provocative", "suggestive", "flirty"
}

NSFW_PHRASES = {
    "spreading legs", "removing bra", "pulling panties", "sucking dick", "licking pussy",
    "penetrating", "fucking scene", "hard cock", "wet pussy", "big tits", "exposed breasts",
    "nipples visible", "ass spread", "thigh gap", "camel toe", "pussy lips", "cum on face",
    "blowjob scene", "anal sex", "titty fuck", "gang rape", "group sex", "public sex",
    "hidden camera", "peeing girl", "pooping girl", "covered in blood", "cutting flesh",
    "snorting cocaine", "injecting heroin", "hallucinating", "smoking weed"
}

SENSITIVE_CONTEXT = {
    "spread", "removing", "pulling", "sucking", "licking", "penetrating",
    "fucking", "hard", "wet", "exposed", "visible", "ass", "tight", "revealing"
}

# Tối ưu hóa hàm NSFW detection
@lru_cache(maxsize=1024)
def detect_nsfw(prompt: str):
    prompt_lower = prompt.lower()
    
    # Kiểm tra từ khóa cao
    words = set(re.findall(r'\b\w+\b', prompt_lower))
    if words & NSFW_HIGH:
        return True
    
    # Kiểm tra cụm từ NSFW
    for phrase in NSFW_PHRASES:
        if phrase in prompt_lower:
            return True
    
    # Kiểm tra từ khóa trung bình + ngữ cảnh nhạy cảm
    medium_matches = words & NSFW_MEDIUM
    if medium_matches and (words & SENSITIVE_CONTEXT):
        return True
    
    return False

def infer(
    prompt: str,
    negative_prompt: str,
    seed: int,
    randomize_seed: bool,
    width: int,
    height: int,
    guidance_scale: float,
    num_inference_steps: int,
    progress=gr.Progress(track_tqdm=True),
):
    print("=" * 60)
    now = datetime.now()
    current_time = now.strftime("%H:%M:%S %d/%m/%Y")
    print(f"⏰ Run at: {current_time}")
    print(f"📝 Prompt: {prompt}")
    print(f"🚫 Negative Prompt: {negative_prompt or '[None]'}")
    
    # Tối ưu hóa kiểm tra NSFW
    if detect_nsfw(prompt):
        raise gr.Error("⚠️ Prompt contains NSFW content. Please use a safe prompt.")
    
    if randomize_seed:
        seed = random.randint(0, np.iinfo(np.int32).max)
        print(f"🎲 Random Seed generated: {seed}")
    else:
        print(f"🔢 Using fixed Seed: {seed}")
    
    # Tối ưu hóa kích thước ảnh
    width = (width // 8) * 8
    height = (height // 8) * 8
    print(f"🖼️ Image Size: {width}x{height}")
    print(f"🎯 Guidance Scale: {guidance_scale}")
    print(f"📈 Inference Steps: {num_inference_steps}")
    
    # Tối ưu hóa generator
    generator = torch.Generator("cpu").manual_seed(seed)
    
    # Sử dụng inference_mode để tăng tốc
    with torch.inference_mode():
        result = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
            # Tối ưu hóa batch
            batch_size=1,
            # Tăng chất lượng
            eta=0.0,  # Giảm nhiễu
            use_karras_sigmas=True,  # Cải thiện chất lượng
        )
    
    image = result.images[0]
    return image, seed

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
.donate-button {
    background-color: #FFDD00 !important;
    color: #000000 !important;
}
"""

# Hàm JavaScript để mở trang donate khi nhấn nút
donate_js = """
function() {
    window.open('https://buymeacoffee.com/harry07?status=1', '_blank');
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            gr.Markdown("## Stable Diffusion v1.5 DEMO")
            # Thêm nút Donate với biểu tượng ☕
            donate_btn = gr.Button("☕ Donate", elem_classes="donate-button")
        
        with gr.Row():
            prompt = gr.Text(label="Prompt", placeholder="Enter your prompt", show_label=False)
            run_button = gr.Button("Generate", variant="primary")
        result = gr.Image(label="Generated Image", show_label=False)
        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(label="Negative prompt", placeholder="Enter negative prompt")
            seed = gr.Slider(label="Seed", minimum=0, maximum=np.iinfo(np.int32).max, step=1, value=0)
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            with gr.Row():
                width = gr.Slider(label="Width", minimum=256, maximum=1024, step=8, value=512)
                height = gr.Slider(label="Height", minimum=256, maximum=1024, step=8, value=512)
            with gr.Row():
                guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=15.0, step=0.1, value=7.5)
                num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=35)
        gr.Examples(
            examples=[
                "A fantasy landscape, vivid colors, sunset light",
                "Portrait of a cyberpunk robot girl, neon lighting",
                "An epic sci-fi scene: spaceship battle in space",
            ],
            inputs=[prompt]
        )
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )
    
    # Gán sự kiện click cho nút Donate
    donate_btn.click(None, js=donate_js)

if __name__ == "__main__":
    demo.launch(share=True)