Update app.py

app.py

@@ -9,16 +9,54 @@ import spaces
 import torch
 import random
 from PIL import Image
+import logging
 
 from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
 
+# Enhanced logging configuration
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
 MAX_SEED = np.iinfo(np.int32).max
 
+class GenerationError(Exception):
+    """Custom exception for generation errors"""
+    pass
+
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
 
+# -------------------- NSFW 检测模型加载 --------------------
+try:
+    logger.info("Loading NSFW detector...")
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    from transformers import AutoProcessor, AutoModelForImageClassification
+    nsfw_processor = AutoProcessor.from_pretrained("Falconsai/nsfw_image_detection")
+    nsfw_model = AutoModelForImageClassification.from_pretrained(
+        "Falconsai/nsfw_image_detection"
+    ).to(device)
+    logger.info("NSFW detector loaded successfully.")
+except Exception as e:
+    logger.error(f"Failed to load NSFW detector: {e}")
+    nsfw_model = None
+    nsfw_processor = None
+
+def detect_nsfw(image: Image.Image, threshold: float = 0.5) -> bool:
+    """Returns True if image is NSFW"""
+    inputs = nsfw_processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = nsfw_model(**inputs)
+        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        nsfw_score = probs[0][1].item()  # label 1 = NSFW
+    return nsfw_score > threshold
+
+
 @spaces.GPU
-def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress()):
+def _infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress()):
     """
     Perform image editing using the FLUX.1 Kontext pipeline.
     
@@ -69,30 +107,61 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+
+    try:
+        if input_image:
+            input_image = input_image.convert("RGB")
+            image = pipe(
+                image=input_image, 
+                prompt=prompt,
+                guidance_scale=guidance_scale,
+                width = input_image.size[0],
+                height = input_image.size[1],
+                num_inference_steps=steps,
+                callback_on_step_end=callback_fn,
+                generator=torch.Generator().manual_seed(seed),
+            ).images[0]
+        else:
+            image = pipe(
+                prompt=prompt,
+                guidance_scale=guidance_scale,
+                num_inference_steps=steps,
+                callback_on_step_end=callback_fn,
+                generator=torch.Generator().manual_seed(seed),
+            ).images[0]
+         # NSFW 检测
+        if nsfw_model and nsfw_processor:
+            if detect_nsfw(image):
+                msg = "Generated image contains NSFW content and cannot be displayed. Please modify your prompt and try again."
+                raise Exception(msg)
+            
+        progress(1, desc="Complete")
+        info = {
+            "status": "success"
+        }
+        return image, info, seed, gr.Button(visible=True)
+    except GenerationError as e:
+        error_info = {
+            "error": str(e),
+            "status": "failed",
+        }
+        return None, error_info, None, None
+    except Exception as e:
+        error_info = {
+            "error": str(e),
+            "status": "failed",
+        }
+        return None, error_info, None, None
+
     
-    if input_image:
-        input_image = input_image.convert("RGB")
-        image = pipe(
-            image=input_image, 
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            width = input_image.size[0],
-            height = input_image.size[1],
-            num_inference_steps=steps,
-            callback_on_step_end=callback_fn,
-            generator=torch.Generator().manual_seed(seed),
-        ).images[0]
-    else:
-        image = pipe(
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=steps,
-            callback_on_step_end=callback_fn,
-            generator=torch.Generator().manual_seed(seed),
-        ).images[0]
-
-    progress(1, desc="Complete")
-    return image, seed, gr.Button(visible=True)
+def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress()):
+    # 调用 GPU 函数
+    image, info, seed, reuse_button = _infer(input_image, prompt,seed,randomize_seed,guidance_scale,steps,progress)
+    # 如果出错，抛出异常
+    if info["status"] == "failed":
+        raise gr.Error(info["error"])
+    # 返回图片
+    return image, seed, reuse_button
 
 @spaces.GPU
 def infer_example(input_image, prompt):