Update app.py

app.py

@@ -4,11 +4,12 @@ import numpy as np
 import spaces
 import torch
 import random
-from PIL import Image
+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(
@@ -28,7 +29,22 @@ 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)
+
+def is_image_nsfw(image):
+    try:
+        results = safety_classifier(image)
+        for result in results:
+            if result['label'] == 'nsfw' and result['score'] > 0.8:
+                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
@@ -78,7 +94,6 @@ pipe.load_lora_weights("vafipas663/Qwen-Edit-2509-Upscale-LoRA",
 pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 MAX_SEED = np.iinfo(np.int32).max
 
-# --- نگاشت نام‌های فارسی به نام‌های داخلی مدل ---
 LORA_MAPPING = {
     "تبدیل عکس به انیمه": "anime",
     "تغییر زاویه دید": "multiple-angles",
@@ -90,16 +105,38 @@ LORA_MAPPING = {
     "افزایش کیفیت (Upscale)": "upscale-image"
 }
 
-# --- گزینه‌های نسبت ابعاد ---
-ASPECT_RATIOS = {
+# لیست ابعاد، شامل گزینه شخصی‌سازی
+ASPECT_RATIOS = [
+    "خودکار (پیش‌فرض)",
+    "۱:۱ (مربع - 1024x1024)",
+    "۱۶:۹ (افقی - 1344x768)",
+    "۹:۱۶ (عمودی - 768x1344)",
+    "شخصی‌سازی (Custom)"
+]
+
+ASPECT_RATIOS_MAP = {
     "خودکار (پیش‌فرض)": "Auto",
-    "۱:۱ (مربع)": (1024, 1024),
-    "۱۶:۹ (افقی - لنداسکیپ)": (1344, 768),
-    "۹:۱۶ (عمودی - استوری)": (768, 1344)
+    "۱:۱ (مربع - 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:
@@ -129,6 +166,13 @@ def update_dimensions_on_upload(image):
     
     return new_width, new_height
 
+# تابع برای نمایش/مخفی کردن اسلایدرهای ابعاد
+def toggle_custom_dimensions(value):
+    if value == "شخصی‌سازی (Custom)":
+        return gr.update(visible=True), gr.update(visible=True)
+    else:
+        return gr.update(visible=False), gr.update(visible=False)
+
 @spaces.GPU(duration=30)
 def infer(
     input_image,
@@ -138,15 +182,24 @@ def infer(
     randomize_seed,
     guidance_scale,
     steps,
-    aspect_ratio_selection, # پارامتر جدید
+    aspect_ratio_selection,
+    custom_width,
+    custom_height,
     progress=gr.Progress(track_tqdm=True)
 ):
     if input_image is None:
         raise gr.Error("لطفاً یک تصویر برای ویرایش بارگذاری کنید.")
 
+    # 1. بررسی امنیت تصویر ورودی
+    if is_image_nsfw(input_image):
+        raise gr.Error("تصویر ورودی دارای محتوای نامناسب تشخیص داده شد.")
+
+    # 2. ترجمه و بررسی متن
     english_prompt = translate_prompt(prompt)
+    if not check_text_safety(english_prompt):
+        raise gr.Error("متن درخواست شامل کلمات غیرمجاز است.")
+
     adapter_internal_name = LORA_MAPPING.get(lora_adapter_persian)
-    
     if adapter_internal_name:
         pipe.set_adapters([adapter_internal_name], adapter_weights=[1.0])
     
@@ -154,26 +207,30 @@ def infer(
         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"
+    
+    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")
     
-    # --- منطق تعیین ابعاد ---
-    if aspect_ratio_selection == "خودکار (پیش‌فرض)" or aspect_ratio_selection is None:
+    # --- منطق ابعاد ---
+    selection_value = ASPECT_RATIOS_MAP.get(aspect_ratio_selection)
+    
+    if selection_value == "Custom":
+        # استفاده از ابعاد شخصی‌سازی شده (باید مضرب 8 باشد)
+        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:
-        # دریافت ابعاد از دیکشنری
-        dims = ASPECT_RATIOS.get(aspect_ratio_selection)
-        if dims and isinstance(dims, tuple):
-            width, height = dims
-        else:
-             # fallback to auto
-            width, height = update_dimensions_on_upload(original_image)
+        # استفاده از ابعاد پیش‌فرض (تاپل)
+        width, height = selection_value
 
     result = pipe(
         image=original_image,
         prompt=english_prompt,
-        negative_prompt=negative_prompt,
+        negative_prompt=final_negative_prompt,
         height=height,
         width=width,
         num_inference_steps=steps,
@@ -181,6 +238,9 @@ def infer(
         true_cfg_scale=guidance_scale,
     ).images[0]
 
+    if is_image_nsfw(result):
+        raise gr.Error("تصویر تولید شده حاوی محتوای نامناسب بود و حذف شد.")
+
     return result, seed
 
 @spaces.GPU(duration=30)
@@ -188,8 +248,8 @@ def infer_example(input_image, prompt, lora_adapter):
     input_pil = input_image.convert("RGB")
     guidance_scale = 1.0
     steps = 4
-    # برای نمونه‌ها حالت خودکار را پیش‌فرض قرار می‌دهیم
-    result, seed = infer(input_pil, prompt, lora_adapter, 0, True, guidance_scale, steps, "خودکار (پیش‌فرض)")
+    # برای نمونه‌ها، ابعاد مهم نیست چون روی خودکار است
+    result, seed = infer(input_pil, prompt, lora_adapter, 0, True, guidance_scale, steps, "خودکار (پیش‌فرض)", 1024, 1024)
     return result, seed
 
 
@@ -214,12 +274,12 @@ async (image) => {
 }
 """
 
-# --- تنظیمات HTML (CSS زیبا و مدرن + JS اجبار تم روشن) ---
+# --- تنظیمات HTML (استایل و اسکریپت) ---
 html_code = """
 <style>
 @import url('https://fonts.googleapis.com/css2?family=Vazirmatn:wght@300;400;500;700&display=swap');
 
-/* --- 1. تنظیمات کلی و اجبار تم روشن --- */
+/* --- تنظیمات تم روشن اجباری --- */
 :root, .dark, body, .gradio-container {
     --body-background-fill: #f5f7fa !important;
     --body-text-color: #1f2937 !important;
@@ -240,7 +300,6 @@ body {
     padding: 10px;
 }
 
-/* --- 2. کانتینر اصلی --- */
 #col-container {
     margin: 0 auto;
     max-width: 980px;
@@ -253,7 +312,6 @@ body {
     border: 1px solid rgba(255,255,255,0.8);
 }
 
-/* --- 3. هدر و توضیحات --- */
 #main-title h1 {
     font-size: 2.4em !important;
     text-align: center;
@@ -273,7 +331,6 @@ body {
     line-height: 1.6;
 }
 
-/* --- 4. استایل ورودی‌ها و لیبل‌ها --- */
 .gr-input-label, span.label-wrap, label span {
     font-weight: 700 !important;
     color: #374151 !important;
@@ -302,7 +359,6 @@ textarea:focus, input[type="text"]:focus {
     border-radius: 12px !important;
 }
 
-/* --- 5. دکمه‌ها --- */
 .primary-btn, button.primary {
     background: linear-gradient(135deg, #10b981 0%, #059669 100%) !important;
     border: none !important;
@@ -335,7 +391,7 @@ textarea:focus, input[type="text"]:focus {
     box-shadow: 0 8px 25px rgba(59, 130, 246, 0.45) !important;
 }
 
-/* --- 6. بخش نمونه‌ها --- */
+/* --- استایل بخش نمونه‌ها --- */
 .gradio-container .prose table, 
 .gradio-container table {
     background-color: #ffffff !important;
@@ -372,12 +428,10 @@ textarea:focus, input[type="text"]:focus {
 }
 
 .gradio-container tbody td span, 
-.gradio-container tbody td p,
-.gradio-container tbody td div {
+.gradio-container tbody td p {
     color: #374151 !important;
 }
 
-/* --- 7. مخفی‌سازی‌ها --- */
 footer { display: none !important; }
 .flagging { display: none !important; }
 
@@ -413,7 +467,6 @@ document.addEventListener('DOMContentLoaded', forceLightMode);
 """
 
 with gr.Blocks() as demo:
-    # تزریق HTML (استایل و اسکریپت)
     gr.HTML(html_code)
     
     with gr.Column(elem_id="col-container"):
@@ -440,7 +493,6 @@ with gr.Blocks() as demo:
             with gr.Column():
                 output_image = gr.Image(label="تصویر نهایی", interactive=False, format="png", height=380)
                 
-                # دکمه دانلود اختصاصی برای ارسال به Iframe
                 download_button = gr.Button("📥 دانلود و ذخیره تصویر", variant="secondary", elem_id="download-btn", elem_classes="primary-btn")
                 
                 with gr.Row():
@@ -450,20 +502,39 @@ with gr.Blocks() as demo:
                         value="تبدیل عکس به انیمه"
                     )
                 
-                # --- تنظیمات پیشرفته با قابلیت انتخاب سایز ---
                 with gr.Accordion("تنظیمات پیشرفته", open=False, visible=True):
                     aspect_ratio_selection = gr.Dropdown(
                         label="ابعاد تصویر خروجی",
-                        choices=list(ASPECT_RATIOS.keys()),
+                        choices=ASPECT_RATIOS,
                         value="خودکار (پیش‌فرض)",
                         interactive=True
                     )
+                    
+                    # --- بخش تنظیمات شخصی‌سازی ابعاد ---
+                    with gr.Row():
+                        custom_width = gr.Slider(
+                            label="عرض دلخواه (Width)",
+                            minimum=256, maximum=2048, step=8, value=1024,
+                            visible=False
+                        )
+                        custom_height = gr.Slider(
+                            label="ارتفاع دلخواه (Height)",
+                            minimum=256, maximum=2048, step=8, value=1024,
+                            visible=False
+                        )
+                    
                     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)
         
-        # تعریف نمونه‌ها (Examples)
+        # رویداد تغییر منوی ابعاد
+        aspect_ratio_selection.change(
+            fn=toggle_custom_dimensions,
+            inputs=aspect_ratio_selection,
+            outputs=[custom_width, custom_height]
+        )
+
         gr.Examples(
             examples=[
                 ["examples/1.jpg", "تبدیل به انیمه کن.", "تبدیل عکس به انیمه"],
@@ -485,14 +556,12 @@ with gr.Blocks() as demo:
             label="نمونه‌ها (برای تست کلیک کنید)"
         )
 
-    # اتصال دکمه اجرا
     run_button.click(
         fn=infer,
-        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps, aspect_ratio_selection],
+        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps, aspect_ratio_selection, custom_width, custom_height],
         outputs=[output_image, seed]
     )
 
-    # اتصال دکمه دانلود
     download_button.click(
         fn=None,
         inputs=[output_image],