app.py

import os
import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image, ImageFilter
from typing import Iterable
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes
from deep_translator import GoogleTranslator
from transformers import pipeline

# --- تعریف تم ---
colors.steel_blue = colors.Color(
    name="steel_blue",
    c50="#EBF3F8",
    c100="#D3E5F0",
    c200="#A8CCE1",
    c300="#7DB3D2",
    c400="#529AC3",
    c500="#4682B4",
    c600="#3E72A0",
    c700="#36638C",
    c800="#2E5378",
    c900="#264364",
    c950="#1E3450",
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --- بارگذاری مدل تشخیص محتوای نامناسب (NSFW) ---
print("Loading Safety Checker...")
safety_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=-1)

def is_image_nsfw(image):
    if image is None: return False
    try:
        results = safety_classifier(image)
        for result in results:
            if result['label'] == 'nsfw' and result['score'] > 0.75:
                return True
        return False
    except Exception as e:
        print(f"Safety check error: {e}")
        return False

# --- بارگذاری مدل اصلی ---
from diffusers import FlowMatchEulerDiscreteScheduler
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading pipeline...")
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2509",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "linoyts/Qwen-Image-Edit-Rapid-AIO",
        subfolder='transformer',
        torch_dtype=dtype,
        device_map='cuda'
    ),
    torch_dtype=dtype
).to(device)

# بارگذاری LoRA ها
pipe.load_lora_weights("autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime", weight_name="Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors", adapter_name="anime")
pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multiple-angles", weight_name="镜头转换.safetensors", adapter_name="multiple-angles")
pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Light_restoration", weight_name="移除光影.safetensors", adapter_name="light-restoration")
pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Relight", weight_name="Qwen-Edit-Relight.safetensors", adapter_name="relight")
pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multi-Angle-Lighting", weight_name="多角度灯光-251116.safetensors", adapter_name="multi-angle-lighting")
pipe.load_lora_weights("tlennon-ie/qwen-edit-skin", weight_name="qwen-edit-skin_1.1_000002750.safetensors", adapter_name="edit-skin")
pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509", weight_name="next-scene_lora-v2-3000.safetensors", adapter_name="next-scene")
pipe.load_lora_weights("vafipas663/Qwen-Edit-2509-Upscale-LoRA", weight_name="qwen-edit-enhance_64-v3_000001000.safetensors", adapter_name="upscale-image")

pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
MAX_SEED = np.iinfo(np.int32).max

LORA_MAPPING = {
    "تبدیل عکس به انیمه": "anime",
    "تغییر زاویه دید": "multiple-angles",
    "اصلاح نور و سایه": "light-restoration",
    "نورپردازی مجدد (Relight)": "relight",
    "نورپردازی چند زاویه‌ای": "multi-angle-lighting",
    "روتوش پوست": "edit-skin",
    "صحنه بعدی (سینمایی)": "next-scene",
    "افزایش کیفیت (Upscale)": "upscale-image"
}

ASPECT_RATIOS_LIST = [
    "خودکار (پیش‌فرض)",
    "۱:۱ (مربع - 1024x1024)",
    "۱۶:۹ (افقی - 1344x768)",
    "۹:۱۶ (عمودی - 768x1344)",
    "شخصی‌سازی (Custom)"
]

ASPECT_RATIOS_MAP = {
    "خودکار (پیش‌فرض)": "Auto",
    "۱:۱ (مربع - 1024x1024)": (1024, 1024),
    "۱۶:۹ (افقی - 1344x768)": (1344, 768),
    "۹:۱۶ (عمودی - 768x1344)": (768, 1344),
    "شخصی‌سازی (Custom)": "Custom"
}

BANNED_WORDS = [
    "nude", "naked", "sex", "porn", "undressed", "nsfw", "erotic", "xxx", 
    "breast", "nipple", "genital", "vagina", "penis", "ass", "butt", "sexual",
    "lingerie", "bikini", "swimwear", "underwear", "fetish", "topless",
    "exhibitionism", "hentai", "ecchi", "18+"
]

def check_text_safety(text):
    text_lower = text.lower()
    for word in BANNED_WORDS:
        if f" {word} " in f" {text_lower} ": 
            return False
    return True

def translate_prompt(text):
    if not text:
        return ""
    try:
        translated = GoogleTranslator(source='auto', target='en').translate(text)
        return translated
    except Exception as e:
        print(f"Translation Error: {e}")
        return text

def update_dimensions_on_upload(image):
    if image is None:
        return 1024, 1024
    original_width, original_height = image.size
    if original_width > original_height:
        new_width = 1024
        aspect_ratio = original_height / original_width
        new_height = int(new_width * aspect_ratio)
    else:
        new_height = 1024
        aspect_ratio = original_width / original_height
        new_width = int(new_height * aspect_ratio)
    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8
    return new_width, new_height

def update_sliders_visibility(choice):
    if choice == "شخصی‌سازی (Custom)":
        return gr.update(visible=True), gr.update(visible=True)
    else:
        return gr.update(visible=False), gr.update(visible=False)

def get_error_html(message):
    return f"""
    <div style="background-color: #fee2e2; border: 1px solid #ef4444; color: #b91c1c; padding: 12px; border-radius: 8px; text-align: center; margin-bottom: 10px; font-weight: bold; display: flex; align-items: center; justify-content: center; gap: 8px;">
        <span style="font-size: 1.2em;">⛔</span>
        {message}
    </div>
    """

def get_success_html(message):
    return f"""
    <div style="background-color: #dcfce7; border: 1px solid #22c55e; color: #15803d; padding: 12px; border-radius: 8px; text-align: center; margin-bottom: 10px; font-weight: bold; display: flex; align-items: center; justify-content: center; gap: 8px;">
        <span style="font-size: 1.2em;">✅</span>
        {message}
    </div>
    """

@spaces.GPU(duration=30)
def infer(
    input_image,
    prompt,
    lora_adapter_persian,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    aspect_ratio_selection,
    custom_width,
    custom_height,
    progress=gr.Progress(track_tqdm=True)
):
    if input_image is None:
        return None, seed, get_error_html("لطفاً ابتدا یک تصویر بارگذاری کنید.")

    if is_image_nsfw(input_image):
        return None, seed, get_error_html("تصویر ورودی دارای محتوای نامناسب است و پردازش نمی‌شود.")

    english_prompt = translate_prompt(prompt)
    if not check_text_safety(english_prompt):
        return None, seed, get_error_html("متن درخواست شامل کلمات غیرمجاز یا غیراخلاقی است.")

    adapter_internal_name = LORA_MAPPING.get(lora_adapter_persian)
    if adapter_internal_name:
        pipe.set_adapters([adapter_internal_name], adapter_weights=[1.0])
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    
    safety_negative = "nsfw, nude, naked, porn, sexual, xxx, breast, nipple, genital, vagina, penis, ass, lingerie, bikini, swimwear, underwear, fetish, topless, gore, violence, blood"
    base_negative = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
    final_negative_prompt = f"{safety_negative}, {base_negative}"

    original_image = input_image.convert("RGB")
    
    selection_value = ASPECT_RATIOS_MAP.get(aspect_ratio_selection)
    
    if selection_value == "Custom":
        width = (int(custom_width) // 8) * 8
        height = (int(custom_height) // 8) * 8
    elif selection_value == "Auto" or selection_value is None:
        width, height = update_dimensions_on_upload(original_image)
    else:
        width, height = selection_value

    try:
        result = pipe(
            image=original_image,
            prompt=english_prompt,
            negative_prompt=final_negative_prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            generator=generator,
            true_cfg_scale=guidance_scale,
        ).images[0]

        if is_image_nsfw(result):
            return None, seed, get_error_html("تصویر تولید شده حاوی محتوای نامناسب بود و حذف شد.")
        
        return result, seed, get_success_html("تصویر با موفقیت ویرایش شد.")

    except Exception as e:
        # اگر خطا مربوط به Quota باشد، اجازه می‌دهیم خود سیستم خطا را تولید کند تا JS آن را بگیرد
        # یا اگر Exception معمولی باشد آن را نشان می‌دهیم
        error_str = str(e)
        if "quota" in error_str.lower() or "exceeded" in error_str.lower():
             raise e # اجازه می‌دهیم خطای اصلی بالا برود تا شامل اعداد زمان باشد
        return None, seed, get_error_html(f"خطا در پردازش: {error_str}")

@spaces.GPU(duration=30)
def infer_example(input_image, prompt, lora_adapter):
    res, s, status = infer(input_image, prompt, lora_adapter, 0, True, 1.0, 4, "خودکار (پیش‌فرض)", 1024, 1024)
    return res, s, status

# --- جاوااسکریپت برای دکمه دانلود ---
js_download_func = """
async (image) => {
    if (!image) {
        alert("لطفاً ابتدا تصویر را تولید کنید.");
        return;
    }
    let fileUrl = image.url;
    if (fileUrl && !fileUrl.startsWith('http')) {
       fileUrl = window.location.origin + fileUrl;
    } else if (!fileUrl && image.path) {
       fileUrl = window.location.origin + "/file=" + image.path;
    }
    window.parent.postMessage({
        type: 'DOWNLOAD_REQUEST',
        url: fileUrl
    }, '*');
}
"""

# --- جاوااسکریپت سراسری (تزریق پیام خطای جدید و تم روشن) ---
js_global_content = """
<script>
document.addEventListener('DOMContentLoaded', () => {
    // 1. Force Light Mode
    const forceLight = () => {
        const body = document.querySelector('body');
        if (body) {
            body.classList.remove('dark');
            body.style.backgroundColor = '#f5f7fa';
            body.style.color = '#333333';
        }
        document.querySelectorAll('.dark').forEach(el => el.classList.remove('dark'));
    };
    forceLight();
    setInterval(forceLight, 1000);

    // 2. NEW GPU Quota Replacer (Hourglass Style)
    const observer = new MutationObserver((mutations) => {
        mutations.forEach((mutation) => {
            if (mutation.addedNodes.length) {
                mutation.addedNodes.forEach((node) => {
                    // فقط روی المان‌های مربوط به خطای گرادیو کار کن
                    if (node.nodeType === 1 && (node.classList.contains('toast-body') || node.classList.contains('error'))) {
                        const originalText = node.innerText;
                        
                        // Regex برای پیدا کردن اعداد زمان
                        const regex = /(\d+)s requested vs. (\d+)s left/;
                        const match = originalText.match(regex);
                        
                        // اگر متن خطا مربوط به Quota بود و قبلا ترجمه نشده بود
                        if (match && !node.dataset.translated) {
                            const requested = match[1];
                            const left = match[2];
                            
                            // **ساخت کارت HTML زیبا**
                            const prettyHtml = `
                                <div style="display: flex; align-items: center; gap: 15px; font-family: 'Tahoma', sans-serif; direction: rtl; padding: 10px;">
                                    <div style="font-size: 2.5em; color: #dc3545;">⏳</div>
                                    <div>
                                        <h4 style="margin: 0; color: #5a6268; font-weight: bold;">ظرفیت سرور تکمیل است!</h4>
                                        <p style="margin: 5px 0 0 0; color: #6c757d; font-size: 0.9em;">
                                            سهمیه رایگان GPU شما برای پردازش یک ویدیوی <b>${requested} ثانیه‌ای</b> کافی نیست.
                                        </p>
                                        <div style="background-color: #f8d7da; border: 1px solid #f5c6cb; border-radius: 5px; padding: 5px 8px; margin-top: 10px; font-size: 0.85em;">
                                            اعتبار باقیمانده: <b>${left} ثانیه</b>
                                        </div>
                                    </div>
                                </div>
                            `;
                            
                            // جایگزینی محتوای قدیمی با کارت جدید
                            node.innerHTML = prettyHtml;
                            
                            // جلوگیری از ترجمه مجدد
                            node.dataset.translated = 'true'; 
                        }
                    }
                });
            }
        });
    });

    observer.observe(document.body, { childList: true, subtree: true });
});
</script>
"""

# --- CSS ---
css_code = """
<style>
@import url('https://fonts.googleapis.com/css2?family=Vazirmatn:wght@300;400;500;700&display=swap');

:root, .dark, body, .gradio-container {
    --body-background-fill: #f5f7fa !important;
    --body-text-color: #1f2937 !important;
    --background-fill-primary: #ffffff !important;
    --background-fill-secondary: #f3f4f6 !important;
    --border-color-primary: #e5e7eb !important;
    --block-background-fill: #ffffff !important;
    --block-label-text-color: #374151 !important;
    --block-title-text-color: #111827 !important;
    --input-background-fill: #ffffff !important;
    color-scheme: light !important;
}

body {
    font-family: 'Vazirmatn', sans-serif !important;
    background-color: #f5f7fa !important;
    margin: 0;
    padding: 10px;
}

#col-container {
    margin: 0 auto;
    max-width: 980px;
    direction: rtl;
    text-align: right;
    padding: 30px;
    background: #ffffff !important;
    border-radius: 24px;
    box-shadow: 0 10px 40px -10px rgba(0,0,0,0.08); 
    border: 1px solid rgba(255,255,255,0.8);
}

#main-title h1 {
    font-size: 2.4em !important;
    text-align: center;
    color: #1a202c !important;
    margin-bottom: 15px;
    font-weight: 800;
    background: -webkit-linear-gradient(45deg, #2563eb, #1e40af);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}

#main-description {
    text-align: center;
    font-size: 1.15em;
    color: #4b5563 !important;
    margin-bottom: 40px;
    line-height: 1.6;
}

.gr-input-label, span.label-wrap, label span {
    font-weight: 700 !important;
    color: #374151 !important;
    font-size: 0.95em !important;
    margin-bottom: 8px !important;
}

textarea, input[type="text"] {
    border: 2px solid #e2e8f0 !important;
    border-radius: 12px !important;
    background-color: #ffffff !important;
    color: #111827 !important;
    padding: 12px !important;
    transition: all 0.3s ease;
    font-family: 'Vazirmatn', sans-serif !important;
}

textarea:focus, input[type="text"]:focus {
    border-color: #3b82f6 !important;
    box-shadow: 0 0 0 4px rgba(59, 130, 246, 0.1) !important;
    outline: none;
}

.gr-dropdown {
    background: #ffffff !important;
    border-radius: 12px !important;
}

.primary-btn, button.primary {
    background: linear-gradient(135deg, #10b981 0%, #059669 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 700 !important;
    font-size: 1.1em !important;
    padding: 14px 28px !important;
    border-radius: 14px !important;
    box-shadow: 0 4px 15px rgba(16, 185, 129, 0.3) !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    cursor: pointer !important;
    width: 100%;
    margin-top: 15px;
}

.primary-btn:hover, button.primary:hover {
    transform: translateY(-2px);
    box-shadow: 0 8px 25px rgba(16, 185, 129, 0.45) !important;
}

.primary-btn:active, button.primary:active {
    transform: translateY(1px);
}

#download-btn {
    background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%) !important;
    box-shadow: 0 4px 15px rgba(59, 130, 246, 0.3) !important;
}
#download-btn:hover {
    box-shadow: 0 8px 25px rgba(59, 130, 246, 0.45) !important;
}

.gradio-container .prose table, 
.gradio-container table {
    background-color: #ffffff !important;
    color: #111827 !important;
    border: 1px solid #e5e7eb !important;
    border-radius: 12px !important;
    overflow: hidden !important;
    width: 100% !important;
    margin-top: 20px !important;
}

.gradio-container thead th {
    background-color: #f3f4f6 !important;
    color: #374151 !important;
    font-weight: 700 !important;
    border-bottom: 2px solid #e5e7eb !important;
    padding: 12px !important;
    text-align: right !important;
}

.gradio-container tbody tr {
    background-color: #ffffff !important;
    border-bottom: 1px solid #f3f4f6 !important;
}

.gradio-container tbody tr:hover {
    background-color: #f9fafb !important;
}

.gradio-container tbody td {
    background-color: #ffffff !important;
    color: #374151 !important;
    padding: 10px !important;
}

.gradio-container tbody td span, 
.gradio-container tbody td p {
    color: #374151 !important;
}

footer { display: none !important; }
.flagging { display: none !important; }

/* استایل Toast برای راست چین کردن */
.toast-body {
    direction: rtl !important;
    text-align: right !important;
}

@media (prefers-color-scheme: dark) {
    body, .gradio-container, .prose, table, tr, td, th {
        background-color: #ffffff !important;
        color: #333333 !important;
    }
}
</style>
"""

# ادغام CSS و JS
combined_html = css_code + js_global_content

# استفاده از gr.Blocks بدون پارامتر js
with gr.Blocks() as demo:
    # تزریق کدها به عنوان HTML
    gr.HTML(combined_html)
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# **ویرایشگر هوشمند آلفا**", elem_id="main-title")
        gr.Markdown(
            "با هوش مصنوعی آلفا تصاویر تونو به مدل های مختلف ویرایش کنید.", 
            elem_id="main-description"
        )

        with gr.Row(equal_height=True):
            with gr.Column():
                input_image = gr.Image(label="بارگذاری تصویر", type="pil", height=320)
                
                prompt = gr.Text(
                    label="دستور ویرایش (به فارسی)",
                    show_label=True,
                    placeholder="مثال: تصویر را به سبک انیمه تبدیل کن...",
                    rtl=True,
                    lines=3
                )
                
                status_box = gr.HTML(label="وضعیت")
                
                run_button = gr.Button("✨ شروع پردازش و ساخت تصویر", variant="primary", elem_classes="primary-btn")

            with gr.Column():
                output_image = gr.Image(label="تصویر نهایی", interactive=False, format="png", height=380)
                
                download_button = gr.Button("📥 دانلود و ذخیره تصویر", variant="secondary", elem_id="download-btn", elem_classes="primary-btn")
                
                with gr.Row():
                    lora_adapter = gr.Dropdown(
                        label="انتخاب سبک ویرایش (LoRA)",
                        choices=list(LORA_MAPPING.keys()),
                        value="تبدیل عکس به انیمه"
                    )
                
                with gr.Accordion("تنظیمات پیشرفته", open=False, visible=True):
                    aspect_ratio_selection = gr.Dropdown(
                        label="ابعاد تصویر خروجی",
                        choices=ASPECT_RATIOS_LIST,
                        value="خودکار (پیش‌فرض)",
                        interactive=True
                    )
                    
                    with gr.Row(visible=False) as custom_dims_row:
                        custom_width = gr.Slider(
                            label="عرض دلخواه (Width)",
                            minimum=256, maximum=2048, step=8, value=1024
                        )
                        custom_height = gr.Slider(
                            label="ارتفاع دلخواه (Height)",
                            minimum=256, maximum=2048, step=8, value=1024
                        )
                    
                    seed = gr.Slider(label="دانه تصادفی (Seed)", minimum=0, maximum=MAX_SEED, step=1, value=0)
                    randomize_seed = gr.Checkbox(label="استفاده از Seed تصادفی", value=True)
                    guidance_scale = gr.Slider(label="میزان وفاداری به متن (Guidance Scale)", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
                    steps = gr.Slider(label="تعداد مراحل پردازش (Steps)", minimum=1, maximum=50, step=1, value=4)
        
        # اصلاح تابع نمایش ردیف اسلایدرها
        def toggle_row(choice):
            if choice == "شخصی‌سازی (Custom)":
                return gr.update(visible=True)
            return gr.update(visible=False)

        aspect_ratio_selection.change(
            fn=toggle_row,
            inputs=aspect_ratio_selection,
            outputs=custom_dims_row
        )

        gr.Examples(
            examples=[
                ["examples/1.jpg", "تبدیل به انیمه کن.", "تبدیل عکس به انیمه"],
                ["examples/5.jpg", "سایه‌ها را حذف کن و نورپردازی نرم به تصویر بده.", "اصلاح نور و سایه"],
                ["examples/4.jpg", "از فیلتر ساعت طلایی با پخش نور ملایم استفاده کن.", "نورپردازی مجدد (Relight)"],
                ["examples/2.jpeg", "دوربین را ۴۵ درجه به سمت چپ بچرخان.", "تغییر زاویه دید"],
                ["examples/7.jpg", "منبع نور را از سمت راست عقب قرار بده.", "نورپردازی چند زاویه‌ای"],
                ["examples/10.jpeg", "کیفیت تصویر را افزایش بده (Upscale).", "افزایش کیفیت (Upscale)"],
                ["examples/7.jpg", "منبع نور را از پایین بتابان.", "نورپردازی چند زاویه‌ای"],
                ["examples/2.jpeg", "زاویه دوربین را به نمای بالا گوشه راست تغییر بده.", "تغییر زاویه دید"],
                ["examples/9.jpg", "دوربین کمی به جلو حرکت می‌کند در حالی که نور خورشید از میان ابرها می‌تابد و درخششی نرم اطراف شبح شخصیت در مه ایجاد می‌کند. سبک سینمایی واقعی.", "صحنه بعدی (سینمایی)"],
                ["examples/8.jpg", "جزئیات پوست سوژه را برجسته‌تر و طبیعی‌تر کن.", "روتوش پوست"],
                ["examples/6.jpg", "دوربین را به نمای پایین به بالا تغییر بده.", "تغییر زاویه دید"],
            ],
            inputs=[input_image, prompt, lora_adapter],
            outputs=[output_image, seed, status_box],
            fn=infer_example,
            cache_examples=False,
            label="نمونه‌ها (برای تست کلیک کنید)"
        )

    run_button.click(
        fn=infer,
        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps, aspect_ratio_selection, custom_width, custom_height],
        outputs=[output_image, seed, status_box],
        api_name="predict"
    )

    download_button.click(
        fn=None,
        inputs=[output_image],
        outputs=None,
        js=js_download_func
    )

if __name__ == "__main__":
    demo.queue(max_size=30).launch(show_error=True)