Update app.py

app.py

@@ -15,7 +15,15 @@ from PIL import Image
 import json
 import base64
 from huggingface_hub import InferenceClient
+import logging
 
+# 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__)
 
 subprocess.check_call([sys.executable, "-m", "pip", "install", "spaces==0.43.0"])
 
@@ -78,6 +86,34 @@ pipe = Flux2Pipeline.from_pretrained(
 )
 pipe.to(device)
 
+class GenerationError(Exception):
+    """Custom exception for generation errors"""
+    pass
+
+# -------------------- 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
+
 # Pull pre-compiled Flux2 Transformer blocks from HF hub
 # flash-attn估计库估计更新了，导致冲突了，不使用预编译的了
 # spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/FLUX.2", variant="fa3")
@@ -171,7 +207,7 @@ def get_duration(prompt_embeds, image_list, width, height, num_inference_steps,
     return max(65, num_inference_steps * step_duration + 10)
 
 @spaces.GPU(duration=get_duration)
-def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress()):
+def _generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress()):
     # Move embeddings to GPU only when inside the GPU decorated function
     prompt_embeds = prompt_embeds.to(device)
     
@@ -187,21 +223,42 @@ def generate_image(prompt_embeds, image_list, width, height, num_inference_steps
         progress(progress_value, desc=f"Image generating, {step + 1}/{num_inference_steps} steps")
         return callback_kwargs
         
-    image = pipe(
-        prompt_embeds=prompt_embeds,
-        image=image_list,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-        generator=generator,
-        width=width,
-        height=height,
-        callback_on_step_end=callback_fn,
-    ).images[0]
-
-    path = save_image(image, "./outputs")
-    progress(1, desc="Complete")
-
-    return path
+    try:
+        image = pipe(
+            prompt_embeds=prompt_embeds,
+            image=image_list,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+            width=width,
+            height=height,
+            callback_on_step_end=callback_fn,
+        ).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 the input image or prompt and try again."
+                raise Exception(msg)
+
+        path = save_image(image, "./outputs")
+        progress(1, desc="Complete")
+        info = {
+            "status": "success"
+        }
+        return path, info
+    except GenerationError as e:
+        error_info = {
+            "error": str(e),
+            "status": "failed",
+        }
+        return None, error_info
+    except Exception as e:
+        error_info = {
+            "error": str(e),
+            "status": "failed",
+        }
+        return None, error_info
 
 def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=2.5, prompt_upsampling=False, progress=gr.Progress()):
     
@@ -233,7 +290,7 @@ def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024,
     
     # 3. Image Generation (GPU bound)
     progress(0.3, desc="Waiting for GPU...")
-    image = generate_image(
+    image, info = _generate_image(
         prompt_embeds, 
         image_list, 
         width, 
@@ -243,6 +300,10 @@ def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024,
         seed, 
         progress
     )
+
+    # 如果出错，抛出异常
+    if info["status"] == "failed":
+        raise gr.Error(info["error"])
     
     return image, seed