Update app.py

app.py

@@ -1,20 +1,14 @@
-import os
 import gc
 import gradio as gr
 import numpy as np
 import torch
-import json
 import spaces
 import random
-import config
 import utils
 import logging
-from PIL import Image, PngImagePlugin
-from datetime import datetime
+from PIL import Image
 from diffusers.models import AutoencoderKL
-from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
-import time
-from typing import List, Dict, Tuple, Optional
+from diffusers import  StableDiffusionXLImg2ImgPipeline
 from config import (
     MODEL,
     MIN_IMAGE_SIZE,
@@ -23,7 +17,8 @@ from config import (
     DEFAULT_NEGATIVE_PROMPT,
     scheduler_list,
 )
-import io
+from transformers import AutoProcessor, AutoModelForImageClassification
+
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -62,6 +57,23 @@ else:
     pipe = None
 
 
+# -------------------- NSFW 检测模型加载 --------------------
+try:
+    logger.info("Loading NSFW detector...")
+    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
+# -----------------------------------------------------------
+
+
+
 class GenerationError(Exception):
     """Custom exception for generation errors"""
     pass
@@ -92,11 +104,20 @@ def validate_dimensions(width: int, height: int) -> None:
         raise GenerationError(f"Height must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE}")
 
 
+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
+
 
 progress=gr.Progress()
 
 @spaces.GPU
-def generate(
+def _generate_on_gpu(
     prompt: str,
     negative_prompt: str,
     width: int,
@@ -177,16 +198,33 @@ def generate(
                 callback_on_step_end=callback2
             ).images
         out_img = images[0] 
+
+        # NSFW 检测
+        if nsfw_model and nsfw_processor:
+            if detect_nsfw(out_img):
+                msg = "Generated image contains NSFW content and cannot be displayed. Please modify your prompt and try again."
+                raise Exception(msg)
+                
         path = utils.save_image(out_img, "./outputs")
         logger.info(f"output path: {path}")
         progress(1, desc="Complete")
-        return path
+
+        info = {
+            "status": "success"
+        }
+        return path, info
     except GenerationError as e:
-        logger.warning(f"Generation validation error: {str(e)}")
-        raise gr.Error(str(e))
+        error_info = {
+            "error": str(e),
+            "status": "failed",
+        }
+        return None, error_info
     except Exception as e:
-        logger.exception("Unexpected error during generation")
-        raise gr.Error(f"Generation failed: {str(e)}")
+        error_info = {
+            "error": str(e),
+            "status": "failed",
+        }
+        return None, error_info
     finally:
         # Cleanup
         torch.cuda.empty_cache()
@@ -200,27 +238,43 @@ def generate(
             
         utils.free_memory()
 
+def generate(
+    prompt: str,
+    negative_prompt: str,
+    width: int,
+    height: int,
+    scheduler: str,
+    opt_strength: float,
+    opt_scale: float,
+    seed: int,
+    randomize_seed: bool,
+    guidance_scale: float,
+    num_inference_steps: int,
+):
+    # 调用 GPU 函数
+    image_path, info = _generate_on_gpu(
+        prompt, negative_prompt,
+        width, height,
+        scheduler,
+        opt_strength, opt_scale,
+        seed, randomize_seed,
+        guidance_scale, num_inference_steps,
+    )
 
-   
+    # 如果出错，抛出异常
+    if info["status"] == "failed":
+        raise gr.Error(info["error"])
 
+    # 返回图片路径
+    return image_path
 
 
 title = "# Anime AI Generator"
 description = "Our AI-Powered Anime Generator turns your ideas into breathtaking AI anime art—perfect for art, storytelling, or personal AI anime wallpaper. Experience more at [Anime AI Generator](https://www.animeaigen.com)."
 
 custom_css = """
-#row-container {
-    align-items: stretch;
-}
-#output-image{
-    flex-grow: 1;
-}
-#output-image *{
-    max-height: none !important;
-}
 """
 
-
 with gr.Blocks(css=custom_css).queue() as demo:
     gr.Markdown(title)
     gr.Markdown(description)
@@ -327,4 +381,4 @@ with gr.Blocks(css=custom_css).queue() as demo:
     )  
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.launch()