Update app.py

app.py

@@ -14,6 +14,15 @@ from deep_translator import GoogleTranslator
 from transformers import pipeline
 from diffusers import Flux2Pipeline, Flux2Transformer2DModel
 from datetime import date
+import subprocess
+import sys
+
+# نصب پکیج spaces در صورت نیاز (برای اطمینان)
+try:
+    import spaces
+except ImportError:
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "spaces==0.43.0"])
+    import spaces
 
 # ==========================================
 # 1. تنظیمات و پیکربندی سیستم (Configuration)
@@ -27,6 +36,9 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# مقادیر سیگما مخصوص حالت توربو (8 مرحله)
+TURBO_SIGMAS = [1.0, 0.6509, 0.4374, 0.2932, 0.1893, 0.1108, 0.0495, 0.00031]
+
 # بارگذاری مدل تشخیص محتوای نامناسب (Safety Checker)
 print("Loading Safety Checker...")
 safety_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=-1)
@@ -55,9 +67,9 @@ BANNED_WORDS = [
 ]
 
 # ==========================================
-# 2. بارگذاری مدل FLUX.2
+# 2. بارگذاری مدل FLUX.2 با حالت TURBO
 # ==========================================
-print("Loading FLUX.2 Pipeline...")
+print("Loading FLUX.2 Pipeline with Turbo LoRA...")
 repo_id = "black-forest-labs/FLUX.2-dev"
 
 dit = Flux2Transformer2DModel.from_pretrained(
@@ -72,6 +84,17 @@ pipe = Flux2Pipeline.from_pretrained(
     transformer=dit,
     torch_dtype=torch.bfloat16
 )
+
+# بارگذاری وزن‌های توربو (Turbo LoRA)
+# این بخش باعث می‌شود مدل به جای ۳۰ مرحله، در ۸ مرحله تصویر بسازد
+print("Loading Turbo LoRA weights...")
+pipe.load_lora_weights(
+    "fal/FLUX.2-dev-Turbo", 
+    weight_name="flux.2-turbo-lora.safetensors"
+)
+pipe.fuse_lora()
+pipe.unload_lora_weights() # ادغام وزن‌ها برای سرعت بیشتر
+
 pipe.to(device)
 
 # بهینه‌سازی ZeroGPU
@@ -102,10 +125,8 @@ usage_data_cache = load_usage_data()
 def is_image_nsfw(image):
     if image is None: return False
     try:
-        # اگر ورودی لیست گالری باشد، اولین تصویر را چک کن
         img_to_check = image
         if isinstance(image, list):
-             # هندل کردن فرمت گالری گرادیو
              if len(image) > 0:
                 img_to_check = image[0][0] if isinstance(image[0], tuple) else image[0]
              else:
@@ -213,13 +234,15 @@ def upsample_prompt_logic(prompt, image_list):
         print(f"Upsampling failed: {e}")
         return prompt
 
+# محاسبه زمان برای GPU (برای حالت توربو بسیار کمتر است)
 def get_duration(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
     num_images = 0 if image_list is None else len(image_list)
     step_duration = 1 + 0.8 * num_images
-    return max(65, num_inference_steps * step_duration + 10)
+    # در حالت توربو همیشه ۸ مرحله داریم، پس زمان محاسبه ثابت و کم است
+    return max(30, 8 * step_duration + 10)
 
 # ==========================================
-# 4. تابع اصلی GPU (Inference)
+# 4. تابع اصلی GPU (Inference) - Turbo Optimized
 # ==========================================
 
 @spaces.GPU(duration=get_duration)
@@ -230,14 +253,18 @@ def generate_image(prompt_embeds, image_list, width, height, num_inference_steps
     pipe_kwargs = {
         "prompt_embeds": prompt_embeds,
         "image": image_list,
-        "num_inference_steps": num_inference_steps,
+        # "num_inference_steps": num_inference_steps, # نادیده گرفته می‌شود به نفع توربو
         "guidance_scale": guidance_scale,
         "generator": generator,
         "width": width,
         "height": height,
     }
 
-    if progress: progress(0, desc="Starting generation...")
+    # اعمال تنظیمات اجباری توربو
+    pipe_kwargs["sigmas"] = TURBO_SIGMAS
+    pipe_kwargs["num_inference_steps"] = 8  # همیشه ۸ مرحله برای توربو
+
+    if progress: progress(0, desc="Starting Turbo generation...")
     image = pipe(**pipe_kwargs).images[0]
     return image
 
@@ -286,8 +313,9 @@ def infer(
         progress(0.3, desc="Encoding...")
         prompt_embeds = remote_text_encoder(final_prompt)
 
-        # Generation (GPU)
-        progress(0.4, desc="Generating...")
+        # Generation (GPU) - TURBO MODE
+        progress(0.4, desc="Generating (Turbo)...")
+        # ورودی num_inference_steps از UI گرفته می‌شود اما در تابع generate_image نادیده گرفته می‌شود (روی ۸ قفل است)
         result_image = generate_image(
             prompt_embeds, image_list, width, height,
             num_inference_steps, guidance_scale, seed, progress
@@ -300,7 +328,7 @@ def infer(
         # 6. محاسبه اعتبار ب��قی‌مانده
         user_record = get_user_record(fingerprint)
         remaining = USAGE_LIMIT - user_record["count"] if user_record else 0
-        success_msg = f"تصویر با موفقیت ساخته شد."
+        success_msg = f"تصویر با موفقیت ساخته شد (Turbo)."
         if subscription_status != 'paid':
              success_msg += f" (اعتبار باقی‌مانده امروز: {remaining})"
 
@@ -505,7 +533,6 @@ footer { display: none !important; }
 # 6. ساخت رابط کاربری (Gradio Blocks)
 # ==========================================
 
-# ******** این خط اصلاح شده است ********
 with gr.Blocks() as demo:
     gr.HTML(js_global_content + css_code)
 
@@ -513,8 +540,8 @@ with gr.Blocks() as demo:
     status_box_input = gr.Textbox(elem_id="status_storage", visible=True)
 
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# **ساخت تصویر با FLUX.2 (پیشرفته)**", elem_id="main-title")
-        gr.Markdown("با استفاده از مدل قدرتمند FLUX.2 متن فارسی خود را به تصاویر شگفت‌انگیز تبدیل کنید.", elem_id="main-description")
+        gr.Markdown("# **ساخت تصویر با FLUX.2 Turbo (فوق سریع)**", elem_id="main-title")
+        gr.Markdown("با استفاده از مدل قدرتمند FLUX.2 Turbo متن فارسی خود را در چند ثانیه به تصاویر شگفت‌انگیز تبدیل کنید.", elem_id="main-description")
         gr.HTML('<div id="badge-container"><span id="user-sub-badge"></span></div>')
 
         with gr.Row():
@@ -539,7 +566,7 @@ with gr.Blocks() as demo:
 
                 status_box = gr.HTML(label="وضعیت")
 
-                run_button = gr.Button("✨ ساخت تصویر", variant="primary", elem_classes="primary-btn", elem_id="run-btn", visible=True)
+                run_button = gr.Button("✨ ساخت سریع تصویر (Turbo)", variant="primary", elem_classes="primary-btn", elem_id="run-btn", visible=True)
                 upgrade_button = gr.Button("💎 خرید نسخه نامحدود", variant="primary", elem_classes="upgrade-btn", elem_id="upgrade-btn", visible=False)
 
                 with gr.Accordion("تنظیمات پیشرفته", open=False):
@@ -550,8 +577,9 @@ with gr.Blocks() as demo:
                         width = gr.Slider(label="عرض (Width)", minimum=256, maximum=MAX_IMAGE_SIZE, step=8, value=1024)
                         height = gr.Slider(label="ارتفاع (Height)", minimum=256, maximum=MAX_IMAGE_SIZE, step=8, value=1024)
                     with gr.Row():
-                        num_inference_steps = gr.Slider(label="تعداد مراحل (Steps)", minimum=1, maximum=50, step=1, value=28)
-                        guidance_scale = gr.Slider(label="میزان وفاداری (Guidance)", minimum=1.0, maximum=10.0, step=0.1, value=3.5)
+                        # این اسلایدر در حالت توربو بی تاثیر است و روی ۸ قفل می شود
+                        num_inference_steps = gr.Slider(label="تعداد مراحل (ثابت در حالت Turbo)", minimum=1, maximum=50, step=1, value=8, interactive=False)
+                        guidance_scale = gr.Slider(label="میزان وفاداری (Guidance)", minimum=1.0, maximum=10.0, step=0.1, value=2.5)
 
             with gr.Column():
                 result = gr.Image(label="تصویر نهایی", show_label=True, interactive=False)