Update app.py

app.py

@@ -29,23 +29,82 @@ colors.steel_blue = colors.Color(
 
 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)
+# --- بارگذاری سیستم امنیتی دوگانه (Dual Safety System) ---
+print("Loading Safety Checkers...")
+
+# مدل اول: سریع و استاندارد
+safety_classifier_1 = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=-1)
+
+# مدل دوم: بسیار دقیق و سخت‌گیر (ViT)
+safety_classifier_2 = pipeline("image-classification", model="AdamCodd/vit-base-nsfw-detector", 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:
+        # بررسی با مدل اول
+        results1 = safety_classifier_1(image)
+        for result in results1:
+            if result['label'] == 'nsfw' and result['score'] > 0.5: # حساسیت متوسط
+                print(f"Safety Check 1 Failed: {result['score']}")
+                return True
+
+        # بررسی با مدل دوم (سخت‌گیرانه برای موارد نیمه برهنه)
+        results2 = safety_classifier_2(image)
+        for result in results2:
+            label = result['label'].lower()
+            score = result['score']
+            
+            # اگر مدل تشخیص دهد nsfw است (حتی با احتمال ۳۰ درصد)
+            if label == 'nsfw' and score > 0.3:
+                print(f"Safety Check 2 (NSFW) Failed: {score}")
+                return True
+            
+            # اگر مدل تشخیص دهد تصویر سکسی یا نیمه برهنه است
+            if label in ['sexy', 'porn', 'hentai'] and score > 0.4:
+                print(f"Safety Check 2 (Partial) Failed: {label} - {score}")
                 return True
+
         return False
     except Exception as e:
         print(f"Safety check error: {e}")
-        return False
+        # در صورت خطا در سیستم امنیتی، جانب احتیاط رعایت شود
+        return True
+
+# --- لیست کلمات ممنوعه (انگلیسی) ---
+# این لیست بعد از ترجمه متن کاربر چک می‌شود
+BANNED_WORDS = [
+    "nsfw", "nude", "naked", "sex", "porn", "erotic", "xxx", "18+", "uncensored",
+    "breast", "nipple", "areola", "cleavage", "topless", "open chest",
+    "genital", "vagina", "penis", "dick", "cock", "pussy", "ass", "butt", "anus",
+    "lingerie", "bikini", "swimwear", "underwear", "panties", "bra",
+    "fetish", "bdsm", "bondage", "exhibitionism", "voyeur",
+    "hentai", "ecchi", "ahegao", "paizuri",
+    "undressed", "stripping", "naked body", "exposed skin", "sheer", "see-through",
+    "rape", "violence", "blood", "gore", "sexual"
+]
+
+def check_text_safety(text):
+    """بررسی متن انگلیسی برای کلمات ممنوعه"""
+    if not text: return True
+    text_lower = text.lower()
+    for word in BANNED_WORDS:
+        # بررسی کلمه به صورت مستقل یا چسبیده
+        if word in text_lower: 
+            print(f"Banned word found: {word}")
+            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
 
-# --- بارگذاری مدل اصلی ---
+# --- بارگذاری مدل اصلی ویرایش تصویر ---
 from diffusers import FlowMatchEulerDiscreteScheduler
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
@@ -54,7 +113,7 @@ from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print("Loading pipeline...")
+print("Loading Generation Pipeline...")
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2509",
     transformer=QwenImageTransformer2DModel.from_pretrained(
@@ -106,30 +165,6 @@ ASPECT_RATIOS_MAP = {
     "شخصی‌سازی (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
@@ -146,12 +181,6 @@ def update_dimensions_on_upload(image):
     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;">
@@ -182,12 +211,15 @@ def infer(
     custom_height,
     progress=gr.Progress(track_tqdm=True)
 ):
+    # 1. بررسی وجود تصویر
     if input_image is None:
         return None, seed, get_error_html("لطفاً ابتدا یک تصویر بارگذاری کنید.")
 
+    # 2. بررسی امنیتی تصویر ورودی (با دو مدل)
     if is_image_nsfw(input_image):
         return None, seed, get_error_html("تصویر ورودی دارای محتوای نامناسب است و پردازش نمی‌شود.")
 
+    # 3. ترجمه و بررسی متن
     english_prompt = translate_prompt(prompt)
     if not check_text_safety(english_prompt):
         return None, seed, get_error_html("متن درخواست شامل کلمات غیرمجاز یا غیراخلاقی است.")
@@ -201,7 +233,15 @@ def infer(
 
     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"
+    # 4. تنظیم دستورات منفی (Negative Prompt) سخت‌گیرانه
+    # این کلمات به هوش مصنوعی می‌گویند که چه چیزهایی نسازد
+    safety_negative = (
+        "nsfw, nude, naked, porn, sexual, xxx, breast, nipple, areola, genital, "
+        "vagina, penis, ass, lingerie, bikini, swimwear, underwear, fetish, "
+        "topless, open chest, revealing clothes, cleavage, see-through, sheer, "
+        "gore, violence, blood, navel, midriff, exposed skin, erotic, ecchi, hentai, "
+        "uncensored, stripping"
+    )
     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}"
 
@@ -229,6 +269,7 @@ def infer(
             true_cfg_scale=guidance_scale,
         ).images[0]
 
+        # 5. بررسی امنیتی تصویر خروجی (با دو مدل)
         if is_image_nsfw(result):
             return None, seed, get_error_html("تصویر تولید شده حاوی محتوای نامناسب بود و حذف شد.")