Update app.py

app.py

@@ -5,7 +5,7 @@ from PIL import Image
 import gc
 import warnings
 
-# Silence the transformers deprecation warning about clean_up_tokenization_spaces
+# Tắt cảnh báo transformers về tokenization spaces
 warnings.filterwarnings(
     "ignore",
     category=FutureWarning,
@@ -17,12 +17,13 @@ warnings.filterwarnings(
 # =============================
 
 MODEL_ID = "timbrooks/instruct-pix2pix"
-MAX_IMAGE_SIZE = 512          # Reduced from 768/1080 to avoid OOM on CPU
+MAX_IMAGE_SIZE = 512          # Có thể giảm xuống 384 nếu vẫn OOM
 DEFAULT_STEPS = 20
 DEFAULT_GUIDANCE = 7.5
+MAX_PROMPT_LENGTH = 200       # Tăng lên để hỗ trợ prompt dài hơn (CLIP mặc định 77, nhưng ta xử lý được dài hơn)
 
 # =============================
-# LOAD MODEL (CPU only)
+# LOAD MODEL (CPU only - NO CUDA/OFFLOAD)
 # =============================
 
 print("Loading model...")
@@ -39,12 +40,11 @@ try:
 
     pipe.to("cpu")
 
-    # Safe memory optimizations for CPU
-    pipe.enable_attention_slicing()          # Reduces memory usage during attention
+    # Chỉ dùng cái này để tiết kiệm RAM, an toàn trên CPU
+    pipe.enable_attention_slicing()
 
-    # Optional: try this instead of sequential offload - usually safer on CPU
-    # Comment out if you still get errors or very slow performance
-    pipe.enable_model_cpu_offload()          # Moves parts of model to CPU when not in use
+    # KHÔNG dùng bất kỳ offload nào: enable_model_cpu_offload, sequential_cpu_offload, v.v.
+    # Nếu bạn thấy dòng nào có "offload" hoặc "cuda" thì xóa ngay
 
     print("Model loaded successfully!")
 
@@ -58,21 +58,27 @@ except Exception as e:
 
 def edit_image(image, prompt, negative_prompt, steps, guidance_scale):
     if image is None:
-        return None, "Please upload an image first."
+        return None, "Vui lòng upload ảnh trước nhé!"
 
     try:
-        # Preprocess input image
+        # Resize ảnh để tránh OOM
         image = image.convert("RGB")
-        original_size = image.size
         image.thumbnail((MAX_IMAGE_SIZE, MAX_IMAGE_SIZE))
 
-        print(f"Processing image resized to: {image.size}")
+        print(f"Ảnh đã resize: {image.size}")
 
-        # Clean memory before inference
+        # Dọn RAM trước khi generate
         gc.collect()
 
         generator = torch.Generator(device="cpu").manual_seed(42)
 
+        # Xử lý prompt dài: cắt bớt nếu quá dài (tránh lỗi indexing)
+        if len(prompt.split()) > MAX_PROMPT_LENGTH:
+            prompt = " ".join(prompt.split()[:MAX_PROMPT_LENGTH])
+            status_msg = f"Prompt quá dài, đã cắt còn {MAX_PROMPT_LENGTH} từ."
+        else:
+            status_msg = "Đang xử lý..."
+
         result = pipe(
             prompt=prompt,
             image=image,
@@ -82,84 +88,58 @@ def edit_image(image, prompt, negative_prompt, steps, guidance_scale):
             generator=generator,
         ).images[0]
 
-        # Optional: resize back to original aspect if desired
-        # result = result.resize(original_size, Image.LANCZOS)
-
-        return result, "Success!"
+        return result, status_msg + " Thành công!"
 
     except RuntimeError as e:
         if "out of memory" in str(e).lower():
-            msg = "❌ Out of memory. Try:\n• Smaller image\n• Fewer steps (10–15)\n• Close other apps"
-            return None, msg
-        return None, f"Runtime error: {str(e)}"
+            return None, "❌ Hết RAM. Thử:\n• Ảnh nhỏ hơn\n• Giảm steps xuống 10-15\n• Đóng tab khác"
+        if "CUDA" in str(e):
+            return None, "Lỗi CUDA - Đã cấu hình CPU only, nếu vẫn lỗi hãy báo mình nhé!"
+        return None, f"Lỗi runtime: {str(e)}"
 
     except Exception as e:
-        return None, f"Error: {str(e)}"
+        return None, f"Lỗi: {str(e)}"
 
 
 # =============================
 # GRADIO INTERFACE
 # =============================
 
-css = """
-.gradio-container {font-family: system-ui, sans-serif;}
-.button {background-color: #4f46e5 !important;}
-"""
-
-with gr.Blocks(title="InstructPix2Pix – CPU Edition", css=css) as demo:
+with gr.Blocks(title="InstructPix2Pix - CPU Edition") as demo:
     gr.Markdown(
         """
         # 🖌 InstructPix2Pix (CPU version)
         
-        Upload an image → Write what you want to change → Generate
+        Upload ảnh → Viết hướng dẫn chỉnh sửa (prompt dài cũng ok, mình tự cắt nếu cần) → Generate
         
-        ⚠️ Running on CPU → each edit takes ~40–180 seconds  
-        Smaller images (≤512px) and fewer steps = faster & more stable
+        ⚠️ Chạy CPU nên chậm (~1-3 phút mỗi lần). Ảnh nhỏ + steps ít = nhanh hơn.
         """
     )
 
     with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(
-                label="Upload Image",
-                type="pil",
-                image_mode="RGB",
-                height=300
-            )
-            prompt = gr.Textbox(
-                label="Instruction (what to change)",
-                placeholder="add sunglasses, change background to beach, make hair longer",
-                lines=3
-            )
-            negative = gr.Textbox(
-                label="Negative prompt (what to avoid)",
-                value="blurry, low quality, deformed, bad anatomy, extra limbs, watermark",
-                lines=2
-            )
-
-        with gr.Column():
-            output_image = gr.Image(
-                label="Edited Result",
-                type="pil",
-                height=300
-            )
-            status = gr.Textbox(label="Status", interactive=False)
-
-    with gr.Accordion("Advanced Settings", open=False):
-        steps_slider = gr.Slider(
-            minimum=10, maximum=50,
-            value=DEFAULT_STEPS,
-            step=5,
-            label="Number of steps (lower = faster)"
-        )
-        guidance = gr.Slider(
-            minimum=1.0, maximum=15.0,
-            value=DEFAULT_GUIDANCE,
-            step=0.5,
-            label="Guidance scale (how strongly to follow prompt)"
-        )
-
-    btn = gr.Button("✨ Generate Edit", variant="primary", scale=1)
+        input_image = gr.Image(label="Ảnh gốc", type="pil", image_mode="RGB", height=350)
+        output_image = gr.Image(label="Ảnh sau chỉnh sửa", type="pil", height=350)
+
+    prompt = gr.Textbox(
+        label="Prompt (hướng dẫn chỉnh sửa - có thể dài)",
+        placeholder="remove all clothing completely, keep face hair body proportions unchanged, high fidelity, realistic skin texture",
+        lines=5,
+        max_lines=10
+    )
+
+    negative = gr.Textbox(
+        label="Negative prompt (tránh những thứ này)",
+        value="blurry, low quality, deformed, bad anatomy, extra limbs, watermark, text",
+        lines=2
+    )
+
+    with gr.Accordion("Cài đặt nâng cao", open=False):
+        steps_slider = gr.Slider(10, 50, value=DEFAULT_STEPS, step=5, label="Số bước inference")
+        guidance = gr.Slider(1.0, 15.0, value=DEFAULT_GUIDANCE, step=0.5, label="Guidance scale")
+
+    btn = gr.Button("✨ Generate", variant="primary")
+
+    status = gr.Textbox(label="Trạng thái", interactive=False)
 
     btn.click(
         fn=edit_image,
@@ -168,8 +148,4 @@ with gr.Blocks(title="InstructPix2Pix – CPU Edition", css=css) as demo:
     )
 
 # Launch
-demo.queue(max_size=3).launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-    show_error=True
-)
+demo.queue(max_size=3).launch(server_name="0.0.0.0", server_port=7860)