update app

app.py

@@ -56,8 +56,8 @@ except Exception as e:
 NCII_MODEL_ID = "hfmlsoc/ncii-light-guard-v01"
 NCII_UNSAFE_LABEL = "ncii"
 NCII_THRESHOLD = 0.5
-NCII_BLOCK_MESSAGE = "Your entered prompt is flagged as NCII (non-consensual intimate imagery) and your request will not be processed. Try with safe prompts."
- 
+NCII_BLOCK_MESSAGE = "You entered prompt is NCII (non-consensual intimate imagery) and your request will not be processed. Try with Safe Prompts."
+
 print("Loading NCII safety guard model...")
 try:
     ncii_guard = hf_pipeline("text-classification", model=NCII_MODEL_ID, device=-1)
@@ -65,7 +65,7 @@ try:
 except Exception as e:
     ncii_guard = None
     print(f"Warning: Could not load NCII guard model ({NCII_MODEL_ID}): {e}")
- 
+
 def check_ncii_safety(prompt_text):
     if ncii_guard is None or not prompt_text or not prompt_text.strip():
         return False, "unknown", 0.0
@@ -81,31 +81,14 @@ def check_ncii_safety(prompt_text):
         print(f"NCII guard inference error: {e}")
     return False, "unknown", 0.0
 
+# ── Examples Config ───────────────────────────────────────────────────────────
 EXAMPLES_CONFIG = [
-    {
-        "images": ["examples/1.jpg"],
-        "prompt": "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed 'Fire-Edit' preserving realistic texture and details.",
-    },
-    {
-        "images": ["examples/2.jpg"],
-        "prompt": "Transform the image into a dotted cartoon style.",
-    },
-    {
-        "images": ["examples/3.jpeg"],
-        "prompt": "Convert it to black and white.",
-    },
-    {
-        "images": ["examples/4.jpg", "examples/5.jpg"],
-        "prompt": "Replace her glasses with the new glasses from image 1.",
-    },
-    {
-        "images": ["examples/8.jpg", "examples/9.png"],
-        "prompt": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result.",
-    },
-    {
-        "images": ["examples/10.jpg", "examples/11.png"],
-        "prompt": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result.",
-    },
+    {"images": ["examples/1.jpg"], "prompt": "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed 'Fire-Edit' preserving realistic texture and details."},
+    {"images": ["examples/2.jpg"], "prompt": "Transform the image into a dotted cartoon style."},
+    {"images": ["examples/3.jpeg"], "prompt": "Convert it to black and white."},
+    {"images": ["examples/4.jpg", "examples/5.jpg"], "prompt": "Replace her glasses with the new glasses from image 1."},
+    {"images": ["examples/8.jpg", "examples/9.png"], "prompt": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result."},
+    {"images": ["examples/10.jpg", "examples/11.png"], "prompt": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result."},
 ]
 
 def make_thumb_b64(path, max_dim=220):
@@ -135,6 +118,7 @@ def encode_full_image(path):
         return ""
 
 def build_client_config():
+    """Static config consumed by the frontend: example cards."""
     examples = []
     for i, ex in enumerate(EXAMPLES_CONFIG):
         examples.append({
@@ -204,9 +188,10 @@ def infer(
     guidance_scale: float,
     steps: int,
 ) -> dict:
-    """Edit one or more images with FireRed-Image-Edit + NCII safety guard.
+    """Edit one or more images with FireRed-Image-Edit-1.1.
 
-    Returns {"image": <base64 PNG data URL>, "seed": <seed used>}.
+    Returns {"image": <base64 PNG data URL>, "seed": <seed used>, "status": "success"} 
+    or {"status": "blocked", "message": <warning>} if NCII triggers.
     """
     gc.collect()
     torch.cuda.empty_cache()
@@ -217,16 +202,23 @@ def infer(
     if not prompt or prompt.strip() == "":
         raise gr.Error("Please enter an edit prompt.")
 
-    # NCII safety check
+    # ── NCII safety check ──
     is_unsafe, _, _ = check_ncii_safety(prompt)
     if is_unsafe:
-        raise gr.Error(NCII_BLOCK_MESSAGE)
+        gc.collect()
+        torch.cuda.empty_cache()
+        # Returning a blocked status instead of raising an error
+        # so the frontend can gracefully catch it and display a warning toast.
+        return {"image": "", "seed": seed, "status": "blocked", "message": NCII_BLOCK_MESSAGE}
 
     if randomize_seed:
         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"
+    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"
+    )
     width, height = update_dimensions_on_upload(pil_images[0])
 
     try:
@@ -240,13 +232,14 @@ def infer(
             generator=generator,
             true_cfg_scale=guidance_scale,
         ).images[0]
-        return {"image": pil_to_b64_png(result_image), "seed": seed}
+        return {"image": pil_to_b64_png(result_image), "seed": seed, "status": "success"}
     except Exception as e:
         raise e
     finally:
         gc.collect()
         torch.cuda.empty_cache()
 
+
 @app.api(name="load_example", queue=False)
 def load_example(idx: float) -> dict:
     """Return base64-encoded example images + prompt for a given example index."""
@@ -265,16 +258,19 @@ def load_example(idx: float) -> dict:
             names.append(os.path.basename(path))
     return {"images": b64_list, "prompt": ex["prompt"], "names": names, "status": "ok"}
 
+
 @app.get("/api/config")
 def client_config():
     """Plain FastAPI route: example card data for the frontend."""
     return CLIENT_CONFIG
 
+
 @app.get("/", response_class=HTMLResponse)
 async def homepage():
     html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
     with open(html_path, "r", encoding="utf-8") as f:
         return f.read()
 
+
 if __name__ == "__main__":
     app.launch(show_error=True, mcp_server=True)