Qwen-Image-Edit-2509-LoRAs-Fast-Fusion

Running on Zero

File size: 33,682 Bytes

import os
import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image
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.orange_red = colors.Color(
    name="orange_red",
    c50="#FFF0E5",
    c100="#FFE0CC",
    c200="#FFC299",
    c300="#FFA366",
    c400="#FF8533",
    c500="#FF4500",
    c600="#E63E00",
    c700="#CC3700",
    c800="#B33000",
    c900="#992900",
    c950="#802200",
)

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

def is_image_nsfw(image):
    if image is None or safety_classifier 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

# --- توابع کمکی متن ---
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

# --- تنظیمات HTML/JS برای مدیریت خطا (نسخه اصلاح شده و قوی) ---
# این بخش دقیقاً همان لاجیک برنامه قبلی برای حذف خطا و نمایش پنجره راهنماست
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. RETRY FUNCTION
    window.retryGeneration = function() {
        const modal = document.getElementById('custom-quota-modal');
        if (modal) modal.remove();

        const runBtn = document.getElementById('run-btn');
        if(runBtn) runBtn.click();
    };
    
    // Close function
    window.closeErrorModal = function() {
         const modal = document.getElementById('custom-quota-modal');
         if (modal) modal.remove();
    };

    // 3. SHOW MODAL FUNCTION
    const showQuotaModal = () => {
        if (document.getElementById('custom-quota-modal')) return;

        const modalHtml = `
            <div id="custom-quota-modal" style="position: fixed; top: 0; left: 0; width: 100%; height: 100%; background: rgba(0,0,0,0.6); backdrop-filter: blur(5px); z-index: 99999; display: flex; align-items: center; justify-content: center; font-family: 'Vazirmatn', sans-serif;">
                <div class="ip-reset-guide-container" style="direction: rtl; text-align: right; max-width: 450px; width: 90%; background: white; border-radius: 20px; padding: 25px; box-shadow: 0 10px 40px rgba(0,0,0,0.3); animation: slideInUp 0.4s ease;">
                    <div class="guide-header">
                        <svg class="guide-header-icon" viewbox="0 0 100 100" xmlns="http://www.w3.org/2000/svg"> 
                            <defs><lineargradient id="grad1" x1="0%" y1="0%" x2="100%" y2="100%"><stop offset="0%" style="stop-color: #667eea; stop-opacity: 1;"></stop><stop offset="100%" style="stop-color: #764ba2; stop-opacity: 1;"></stop></lineargradient></defs> 
                            <circle cx="50" cy="50" r="45" fill="url(#grad1)" opacity="0.1"></circle>
                            <circle cx="50" cy="50" r="35" fill="none" stroke="url(#grad1)" stroke-width="2" opacity="0.3"></circle>
                            <path d="M35 50 L45 60 L65 40" stroke="url(#grad1)" stroke-width="4" fill="none" stroke-linecap="round" stroke-linejoin="round"></path>
                            <circle cx="65" cy="35" r="8" fill="#fee140"></circle>
                            <path d="M62 35 L68 35 M65 32 L65 38" stroke="white" stroke-width="2" stroke-linecap="round"></path> 
                        </svg>
                        <div>
                            <h2>یک قدم تا ساخت تصاویر جدید</h2>
                            <p>نیازمند تغییر نقطه دستیابی</p>
                        </div>
                    </div>
                    
                    <div class="guide-content">
                        <div class="info-card">
                            <div class="info-card-header">
                                <svg class="info-card-icon" viewbox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M13 2L3 14h9l-1 8 10-12h-9l1-8z" fill="#667eea" opacity="0.2"></path><path d="M13 2L3 14h9l-1 8 10-12h-9l1-8z" stroke="#667eea" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path></svg> 
                                <span class="info-card-title">راه حل سریع</span>
                            </div>
                            <p>طبق ویدیو آموزشی پایین بین نقطه دستیابی جابجا شوید تلاش مجدد بزنید تا تصاویر مجدداً تولید بشه.</p>
                        </div>
                        
                        <div class="summary-section">
                            <div class="summary-header">
                                <svg class="summary-icon" viewbox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><circle cx="12" cy="12" r="10" fill="#56ab2f" opacity="0.2"></circle><path d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2z" stroke="#56ab2f" stroke-width="2"></path><path d="M9 12l2 2 4-4" stroke="#56ab2f" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path></svg> 
                                <span class="summary-title">خلاصه راهنما</span>
                            </div>
                            <div class="summary-text">هربار که این صفحه را مشاهده کردید: از اینترنت سیم‌کارت استفاده کنید، VPN را خاموش کرده و طبق ویدیو آموزشی پایین نقطه دستیابی رو تغییر دهید. «تلاش مجدد» کلیک کنید. با این روش ساده می‌توانید به صورت نامحدود تصویر بسازید! ☘️</div>
                        </div>
                        
                        <div class="video-button-container">
                            <button onclick="parent.postMessage({ type: 'NAVIGATE_TO_URL', url: '#/nav/online/news/getSingle/1149635/eyJpdiI6IjhHVGhPQWJwb3E0cjRXbnFWTW5BaUE9PSIsInZhbHVlIjoiS1V0dTdvT21wbXAwSXZaK1RCTG1pVXZqdlFJa1hXV1RKa2FLem9zU3pXMjd5MmlVOGc2YWY0NVdNR3h3Smp1aSIsIm1hYyI6IjY1NTA5ZDYzMjAzMTJhMGQyMWQ4NjA4ZDgyNGZjZDVlY2MyNjdiMjA2NWYzOWRjY2M4ZmVjYWRlMWNlMWQ3ODEiLCJ0YWciOiIifQ==/21135210' }, '*')" class="elegant-video-button"> 
                                <svg class="elegant-video-button-icon" xmlns="http://www.w3.org/2000/svg" viewbox="0 0 24 24"><path d="M8 5v14l11-7z"></path></svg> 
                                <span>دیدن ویدیو آموزشی استفاده نامحدود</span> 
                            </button>
                        </div>
                    </div>
                    
                    <div class="guide-actions">
                        <button class="action-button back-button" onclick="window.closeErrorModal()"> 
                            <svg class="action-button-icon" viewbox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M19 12H5M12 19l-7-7 7-7" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path></svg> 
                            <span>بازگشت</span> 
                        </button>
                        <button class="action-button retry-button" onclick="window.retryGeneration()"> 
                            <svg class="action-button-icon" viewbox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M23 4v6h-6M1 20v-6h6" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path><path d="M20.49 9A9 9 0 0 0 5.64 5.64L1 10m22 4l-4.64 4.36A9 9 0 0 1 3.51 15" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path></svg> 
                            <span>تلاش مجدد</span> 
                        </button>
                    </div>
                </div>
            </div>
        `;
        
        document.body.insertAdjacentHTML('beforeend', modalHtml);
    };

    // 4. SCANNER
    setInterval(() => {
        const potentialErrors = document.querySelectorAll('.toast-body, .error, .toast-wrap, .eta-bar, div[class*="error"]');
        
        potentialErrors.forEach(el => {
            const text = el.innerText || "";
            if (text.toLowerCase().includes('quota') || text.toLowerCase().includes('exceeded')) {
                 
                 showQuotaModal();
                 
                 // Immediately hide the Gradio error
                 el.style.display = 'none';
                 el.style.opacity = '0';
                 el.innerText = '';
                 
                 const parentWrap = el.closest('.toast-wrap');
                 if(parentWrap) parentWrap.style.display = 'none';
            }
        });
    }, 100);
});
</script>
"""

# --- CSS Updated ---
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;
}

/* --- IP Reset Guide CSS --- */
:root {
    --guide-bg: rgba(255, 255, 255, 0.98);
    --guide-border: rgba(102, 126, 234, 0.2);
    --guide-text-title: #2d3748;
    --guide-text-body: #4a5568;
    --guide-accent: #667eea;
    --primary-gradient-guide: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    --success-gradient-guide: linear-gradient(135deg, #56ab2f 0%, #a8e063 100%);
    --radius-md-guide: 12px;
    --radius-lg-guide: 20px;
    --shadow-sm: 0 1px 2px 0 rgba(26, 32, 44, 0.03);
    --shadow-md: 0 4px 6px -1px rgba(26, 32, 44, 0.05), 0 2px 4px -2px rgba(26, 32, 44, 0.04);
    --shadow-xl: 0 20px 25px -5px rgba(26, 32, 44, 0.07), 0 8px 10px -6px rgba(26, 32, 44, 0.05);
}

@keyframes float {
    0%, 100% { transform: translateY(0px); }
    50% { transform: translateY(-10px); }
}
@keyframes slideInUp {
    from { opacity: 0; transform: translateY(30px); }
    to { opacity: 1; transform: translateY(0); }
}

.guide-header { display: flex; align-items: center; margin-bottom: 20px; }
.guide-header-icon { width: 50px; height: 50px; margin-left: 15px; animation: float 3s ease-in-out infinite; flex-shrink: 0; }
.guide-header h2 { font-size: 1.3rem; color: var(--guide-text-title); font-weight: 700; margin: 0; }
.guide-header p { color: var(--guide-text-body); font-size: 0.8rem; margin-top: 5px; margin-bottom: 0; }

.guide-content { font-size: 0.9rem; color: var(--guide-text-body); line-height: 1.6; }

.info-card { background: linear-gradient(135deg, #667eea15 0%, #764ba215 100%); border: 1px solid rgba(102, 126, 234, 0.2); border-radius: var(--radius-md-guide); padding: 15px; margin: 15px 0; position: relative; overflow: hidden; }
.info-card p { font-size: 0.85rem; line-height: 1.6; margin: 0; }
.info-card::before { content: ''; position: absolute; top: 0; right: 0; width: 4px; height: 100%; background: var(--primary-gradient-guide); }
.info-card-header { display: flex; align-items: center; margin-bottom: 10px; }
.info-card-icon { width: 20px; height: 20px; margin-left: 10px; }
.info-card-title { font-weight: 600; color: var(--guide-text-title); font-size: 0.95rem; }

.summary-section { margin-top: 15px; padding: 15px; border-radius: var(--radius-md-guide); background: linear-gradient(135deg, #56ab2f15 0%, #a8e06315 100%); border: 1px solid rgba(86, 171, 47, 0.2); position: relative; overflow: hidden; }
.summary-section::before { content: ''; position: absolute; top: 0; right: 0; width: 4px; height: 100%; background: var(--success-gradient-guide); }
.summary-header { display: flex; align-items: center; margin-bottom: 8px; }
.summary-icon { width: 20px; height: 20px; margin-left: 10px; }
.summary-title { font-weight: 600; color: #2f5a33; font-size: 0.95rem; }
.summary-text { color: #2f5a33; font-size: 0.85rem; line-height: 1.6; }

/* Tutorial Button */
.video-button-container { text-align: center; margin: 25px 0 15px 0; width: 100%; }
.elegant-video-button { 
    display: inline-flex !important; 
    align-items: center; 
    justify-content: center; 
    padding: 10px 24px !important; 
    background-color: #fff !important; 
    color: var(--guide-accent) !important; 
    border: 1px solid #e2e8f0 !important; 
    text-decoration: none; 
    border-radius: 50px !important; 
    font-weight: 600 !important; 
    font-size: 0.9rem !important; 
    cursor: pointer !important; 
    font-family: inherit; 
    transition: all 0.3s ease !important; 
    box-shadow: 0 2px 10px rgba(0,0,0,0.05) !important; 
    width: auto !important;
}
.elegant-video-button:hover { 
    background: var(--primary-gradient-guide) !important; 
    color: white !important; 
    border-color: transparent !important; 
    transform: translateY(-2px); 
    box-shadow: 0 6px 16px rgba(102, 126, 234, 0.3) !important; 
}
.elegant-video-button-icon { width: 18px; height: 18px; margin-left: 8px; fill: currentColor; }

/* Action Buttons */
.guide-actions { 
    display: flex !important; 
    gap: 12px !important; 
    margin-top: 20px; 
    padding-top: 20px; 
    border-top: 1px solid #e2e8f0; 
    width: 100% !important;
}
.action-button { 
    padding: 12px 15px !important; 
    border: none !important; 
    border-radius: 12px !important; 
    font-size: 0.95rem !important; 
    font-weight: 600 !important; 
    cursor: pointer !important; 
    flex: 1 !important; 
    transition: all 0.3s ease !important; 
    display: flex !important; 
    align-items: center; 
    justify-content: center; 
    font-family: inherit; 
    height: 48px !important;
}
.action-button-icon { width: 20px; height: 20px; margin-right: 0; margin-left: 8px; }

.back-button { 
    background: white !important; 
    color: var(--guide-text-body) !important; 
    border: 2px solid #e2e8f0 !important; 
}
.back-button:hover { 
    background: #f7fafc !important; 
    border-color: var(--guide-accent) !important; 
    transform: translateY(-2px); 
    box-shadow: var(--shadow-md) !important; 
}

.retry-button { 
    background: var(--primary-gradient-guide) !important; 
    color: white !important; 
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3) !important; 
}
.retry-button:hover { 
    transform: translateY(-2px); 
    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.4) !important; 
}

/* --- Main App CSS --- */
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; }

/* Force toast transparency */
.toast-body {
    direction: rtl !important;
    text-align: right !important;
    background: transparent !important;
    box-shadow: none !important;
    border: none !important;
    padding: 0 !important;
    max-width: 100% !important;
    width: auto !important;
}
.toast-wrap {
    background: transparent !important;
    border: none !important;
    box-shadow: none !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

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

print("Loading Qwen Image Edit 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)

print("Loading and Fusing Lightning LoRA...")
pipe.load_lora_weights("lightx2v/Qwen-Image-Lightning", weight_name="Qwen-Image-Lightning-4steps-V2.0-bf16.safetensors", adapter_name="lightning")
pipe.fuse_lora(adapter_names=["lightning"], lora_scale=1.0)

print("Loading Task Adapters...")
pipe.load_lora_weights("tarn59/apply_texture_qwen_image_edit_2509", weight_name="apply_texture_v2_qwen_image_edit_2509.safetensors", adapter_name="texture")
pipe.load_lora_weights("ostris/qwen_image_edit_inpainting", weight_name="qwen_image_edit_inpainting.safetensors", adapter_name="fusion")
pipe.load_lora_weights("ostris/qwen_image_edit_2509_shirt_design", weight_name="qwen_image_edit_2509_shirt_design.safetensors", adapter_name="shirt_design")
pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Fusion", weight_name="溶图.safetensors", adapter_name="fusion-x")
pipe.load_lora_weights("oumoumad/Qwen-Edit-2509-Material-transfer", weight_name="material-transfer_000004769.safetensors", adapter_name="material-transfer")
pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Light-Migration", weight_name="参考色调.safetensors", adapter_name="light-migration")

try:
    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
    print("Flash Attention 3 Processor set successfully.")
except Exception as e:
    print(f"Could not set FA3 processor (likely hardware mismatch): {e}. using default attention.")

MAX_SEED = np.iinfo(np.int32).max

LORA_MAPPING_PERSIAN = {
    "ویرایش بافت (Texture)": "Texture Edit",
    "طراحی لباس (Shirt Design)": "Cloth-Design-Fuse",
    "ترکیب اشیاء (Fusion)": "Fuse-Objects",
    "ترکیب پیشرفته (Super Fusion)": "Super-Fusion",
    "انتقال نور (Light Migration)": "Light-Migration",
    "انتقال متریال (Material Transfer)": "Material-Transfer"
}

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 // 16) * 16
    new_height = (new_height // 16) * 16
    return new_width, new_height

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; direction: rtl;">
        <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; direction: rtl;">
        <span style="font-size: 1.2em;">✅</span>
        {message}
    </div>
    """

@spaces.GPU(duration=30)
def infer(
    image_1,
    image_2,
    prompt,
    lora_adapter_persian,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    progress=gr.Progress(track_tqdm=True)
):
    if image_1 is None or image_2 is None:
        return None, seed, get_error_html("لطفاً هر دو تصویر (تصویر پایه و تصویر مرجع) را بارگذاری کنید.")

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

    lora_adapter = LORA_MAPPING_PERSIAN.get(lora_adapter_persian, "Texture Edit")
    
    english_prompt = translate_prompt(prompt)
    if not english_prompt:
        # Default prompts based on internal logic
        if lora_adapter == "Cloth-Design-Fuse": english_prompt = "Put this design on their shirt."
        elif lora_adapter == "Texture Edit": english_prompt = "Apply texture to object."
        elif lora_adapter == "Fuse-Objects": english_prompt = "Fuse object into background."
        elif lora_adapter == "Super-Fusion": english_prompt = "Blend the product into the background, correct its perspective and lighting, and make it naturally integrated with the scene."
        elif lora_adapter == "Material-Transfer": english_prompt = "change materials of image1 to match the reference in image2"
        elif lora_adapter == "Light-Migration": english_prompt = "Refer to the color tone, remove the original lighting from Image 1, and relight Image 1 based on the lighting and color tone of Image 2."

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

    adapters_map = {
        "Texture Edit": "texture",
        "Fuse-Objects": "fusion",
        "Cloth-Design-Fuse": "shirt_design",
        "Super-Fusion": "fusion-x",
        "Material-Transfer": "material-transfer",
        "Light-Migration": "light-migration",
    }
    
    active_adapter = adapters_map.get(lora_adapter)
    if active_adapter:
        pipe.set_adapters([active_adapter], adapter_weights=[1.0])
    else:
        pipe.set_adapters([], adapter_weights=[])

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    negative_prompt = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry, nsfw, nude"

    img1_pil = image_1.convert("RGB")
    img2_pil = image_2.convert("RGB")
    
    width, height = update_dimensions_on_upload(img1_pil)

    try:
        result = pipe(
            image=[img1_pil, img2_pil],
            prompt=english_prompt,
            negative_prompt=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:
        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(image_1, image_2, prompt, lora_adapter_persian):
    res, s, status = infer(image_1, image_2, prompt, lora_adapter_persian, 0, True, 1.0, 4)
    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
    }, '*');
}
"""

with gr.Blocks() as demo:
    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(scale=1):
                with gr.Row():
                    image_1 = gr.Image(label="۱. تصویر پایه (اصلی)", type="pil", height=290)
                    image_2 = gr.Image(label="۲. تصویر مرجع (استایل/بافت)", type="pil", height=290)

                prompt = gr.Text(
                    label="دستور ویرایش (به فارسی)",
                    show_label=True,
                    placeholder="مثال: بافت چوب را روی لیوان اعمال کن...",
                    rtl=True,
                    lines=3
                )
                
                status_box = gr.HTML(label="وضعیت")
                
                # IMPORTANT: elem_id="run-btn" for JS
                run_button = gr.Button("✨ شروع پردازش و ترکیب تصاویر", variant="primary", elem_classes="primary-btn", elem_id="run-btn")

                with gr.Accordion("تنظیمات پیشرفته", open=False, visible=True):
                    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)

            with gr.Column(scale=1):
                output_image = gr.Image(label="تصویر نهایی", interactive=False, format="png", height=350)
                download_button = gr.Button("📥 دانلود و ذخیره تصویر", variant="secondary", elem_id="download-btn", elem_classes="primary-btn")
                
                with gr.Row():
                    lora_adapter = gr.Dropdown(
                        label="انتخاب سبک ویرایش",
                        choices=list(LORA_MAPPING_PERSIAN.keys()),
                        value="ویرایش بافت (Texture)",
                    )        
 
        gr.Examples(
            examples=[
                ["examples/M1.jpg", "examples/M2.jpg", "با توجه به تن رنگ، نورپردازی اصلی تصویر ۱ را حذف کن و بر اساس نور و رنگ تصویر ۲ مجدداً نورپردازی کن.", "انتقال نور (Light Migration)"],
                ["examples/Cloth2.jpg", "examples/Design2.png", "این طرح را روی پیراهن قرار بده.", "طراحی لباس (Shirt Design)"],
                ["examples/Cup1.png", "examples/Wood1.png", "بافت چوب را روی لیوان اعمال کن.", "ویرایش بافت (Texture)"],
                ["examples/Cloth1.jpg", "examples/Design1.png", "این طرح را روی پیراهن قرار بده.", "طراحی لباس (Shirt Design)"],
                ["examples/F3.jpg", "examples/F4.jpg", "عینک او را با عینک جدید از تصویر ۱ جایگزین کن.", "ترکیب پیشرفته (Super Fusion)"],
                ["examples/Chair.jpg", "examples/Material.jpg", "متریال تصویر ۱ را تغییر بده تا با مرجع تصویر ۲ مطابقت داشته باشد.", "انتقال متریال (Material Transfer)"],
                ["examples/F1.jpg", "examples/F2.jpg", "بطری کوچک را روی میز قرار بده.", "ترکیب پیشرفته (Super Fusion)"],
                ["examples/Mug1.jpg", "examples/Texture1.jpg", "طرح تصویر ۲ را روی لیوان اعمال کن.", "ویرایش بافت (Texture)"],
                ["examples/Cat1.jpg", "examples/Glass1.webp", "یک گربه که عینک زده است.", "ترکیب اشیاء (Fusion)"],
            ],
            inputs=[image_1, image_2, prompt, lora_adapter],
            outputs=[output_image, seed, status_box],
            fn=infer_example,
            cache_examples=False,
            label="نمونه‌ها (برای تست کلیک کنید)"
        )

    run_button.click(
        fn=infer,
        inputs=[image_1, image_2, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
        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=50).launch(show_error=True)