Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import spaces
 import torch
 import random
+import gc
 from PIL import Image
 from typing import Iterable
 from gradio.themes import Soft
@@ -185,11 +186,17 @@ def translate_prompt(text):
         return text
 
 # ==========================================
-# 4. بارگذاری مدل و LoRA ها
+# 4. بارگذاری مدل و LoRA ها + مدیریت حافظه
 # ==========================================
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dtype = torch.bfloat16
 
+# تابع مهم برای خالی کردن حافظه و جلوگیری از ارور NVML
+def flush():
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+
 from diffusers import FlowMatchEulerDiscreteScheduler
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
@@ -319,6 +326,9 @@ def infer(
     subscription_status,
     progress=gr.Progress(track_tqdm=True)
 ):
+    # پاکسازی حافظه قبل از شروع
+    flush()
+
     # 1. بررسی اعتبار
     if subscription_status != 'paid':
         user_record = get_user_record(fingerprint)
@@ -340,7 +350,6 @@ def infer(
     if not check_text_safety(english_prompt):
         return None, seed, get_error_html("متن درخواست شامل کلمات غیرمجاز است."), gr.update(visible=True), gr.update(visible=False)
 
-    # اگر پرامپت خالی بود، پرامپت پیش‌فرض انگلیسی تنظیم می‌شود (نیاز به ترجمه ندارد)
     if not english_prompt:
         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."
@@ -349,7 +358,7 @@ def infer(
         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."
 
-    # 5. کسر اعتبار (اگر کاربر رایگان است)
+    # 5. کسر اعتبار
     if subscription_status != 'paid':
         consume_quota(fingerprint)
 
@@ -419,13 +428,17 @@ def infer(
     except Exception as e:
         error_str = str(e)
         if "quota" in error_str.lower() or "exceeded" in error_str.lower():
-             raise e # Raise for JS Scanner
+             raise e 
         return None, seed, get_error_html(f"خطا در پردازش: {error_str}"), gr.update(visible=True), gr.update(visible=False)
+    finally:
+        # پاکسازی حافظه بعد از پایان (موفق یا ناموفق)
+        flush()
 
 @spaces.GPU(duration=30)
 def infer_example(image_1, image_2, prompt, lora_adapter):
-    # مثال‌ها: ابعاد خودکار، اعتبار رایگان
+    flush() # پاکسازی قبل از مثال
     res, s, status, btn1, btn2 = infer(image_1, image_2, prompt, lora_adapter, 0, True, 1.0, 4, "خودکار (پیش‌فرض)", 1024, 1024, "example", "paid")
+    flush() # پاکسازی بعد از مثال
     return res, s
 
 # ==========================================
@@ -443,7 +456,7 @@ async (image) => {
 
 js_upgrade_func = """() => { window.parent.postMessage({ type: 'NAVIGATE_TO_PREMIUM' }, '*'); }"""
 
-# جاوا اسکریپت با مودال دقیق و استایل دار
+# جاوا اسکریپت اصلاح شده برای رفع باگ دکمه‌ها
 js_global = """
 <script>
 document.addEventListener('DOMContentLoaded', () => {
@@ -512,21 +525,40 @@ document.addEventListener('DOMContentLoaded', () => {
     init();
 
     // ---------------------------------------------
-    // GPU QUOTA HANDLER (Exact Styled Modal)
+    // GPU QUOTA HANDLER & UI FIXES
     // ---------------------------------------------
+    
+    // پرچم برای جلوگیری از باز شدن مجدد مودال
+    window.quotaModalActive = false;
+
     window.retryGeneration = function() {
         const modal = document.getElementById('custom-quota-modal');
         if (modal) modal.remove();
+        window.quotaModalActive = false; // Reset flag
+        
         const runBtn = document.getElementById('run-btn');
-        if(runBtn) runBtn.click();
+        if(runBtn) {
+            runBtn.click();
+            // گاهی اوقات نیاز به کلیک دوم هست، این کد اطمینان حاصل می‌کند
+            setTimeout(() => runBtn.click(), 100);
+        }
     };
+    
     window.closeErrorModal = function() {
          const modal = document.getElementById('custom-quota-modal');
          if (modal) modal.remove();
+         window.quotaModalActive = false; // Reset flag
+         
+         // پاک کردن متن ارورها برای جلوگیری از تریگر مجدد اسکنر
+         document.querySelectorAll('.toast-body, .error, div[class*="error"]').forEach(el => {
+            el.innerText = '';
+            el.style.display = 'none';
+         });
     };
 
     const showQuotaModal = () => {
-        if (document.getElementById('custom-quota-modal')) return;
+        if (window.quotaModalActive || document.getElementById('custom-quota-modal')) return;
+        window.quotaModalActive = true;
 
         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;">
@@ -586,7 +618,7 @@ document.addEventListener('DOMContentLoaded', () => {
         `;
         
         document.body.insertAdjacentHTML('beforeend', modalHtml);
-        setTimeout(() => { window.closeErrorModal(); }, 10000);
+        // اتوماتیک بسته نمی‌شود تا کاربر دکمه بازگشت را بزند
     };
 
     setInterval(() => {