Update app.py

app.py

@@ -12,6 +12,8 @@ import uvicorn
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 import trimesh
+from transformers import AutoProcessor, AutoModelForImageClassification
+from PIL import Image
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--model_path", type=str, default='tencent/Hunyuan3D-2mini')
@@ -38,6 +40,11 @@ CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 HTML_HEIGHT = 500
 HTML_WIDTH = 500
 
+# -------------------- NSFW 检测模型加载 --------------------
+nsfw_processor = AutoProcessor.from_pretrained("Falconsai/nsfw_image_detection")
+nsfw_model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection").to(args.device)
+# -----------------------------------------------------------
+
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -131,10 +138,22 @@ floater_remove_worker = FloaterRemover()
 degenerate_face_remove_worker = DegenerateFaceRemover()
 face_reduce_worker = FaceReducer()
 
+
+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(args.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(duration=40)
-def gen_shape(
+def _gen_shape_on_gpu(
     image=None,
     steps=50,
     guidance_scale=7.5,
@@ -152,14 +171,28 @@ def gen_shape(
 
 
     if image is None:
-        raise gr.Error("Please provide either a caption or an image.")
+        error_info = {
+            "error": "Please provide either a caption or an image.",
+            "status": "failed",
+        }
+        return None,None,None,None,error_info
 
+    rgbImage = image.convert('RGB')
+    
+    # NSFW 检测
+    if nsfw_model and nsfw_processor:
+        if detect_nsfw(rgbImage):
+            error_info = {
+                "error": "The input image contains NSFW content and cannot be used. Please provide a different image and try again.",
+                "status": "failed",
+            }
+            return None,None,None,None,error_info
 
     seed = int(randomize_seed_fn(seed, randomize_seed))
     octree_resolution = int(octree_resolution)
     save_folder = gen_save_folder()
     # 先移除背景
-    image = rmbg_worker(image.convert('RGB'))
+    image = rmbg_worker(rgbImage)
 
     # 生成模型
     generator = torch.Generator()
@@ -188,7 +221,11 @@ def gen_shape(
         torch.cuda.empty_cache()
     
     if path is None:
-        raise gr.Error('Please generate a mesh first.')
+        error_info = {
+            "error": "'Please generate a mesh first.'",
+            "status": "failed",
+        }
+        return None,None,None,None,error_info
 
     # 简化模型
     print(f'exporting {path}')
@@ -221,9 +258,37 @@ def gen_shape(
 
 
     progress(1,desc="Complete")
-    return model_viewer_html, gr.update(value=sourceObjPath, interactive=True), glbPath, objPath
+    info = {
+        "status": "success"
+    }
+    return model_viewer_html, gr.update(value=sourceObjPath, interactive=True), glbPath, objPath, info
 
 
+def gen_shape(
+    image=None,
+    steps=50,
+    guidance_scale=7.5,
+    seed=1234,
+    octree_resolution=256,
+    num_chunks=200000,
+    target_face_num=10000,
+    randomize_seed: bool = False,
+):
+    # 调用 GPU 函数
+    html_export_mesh,file_export,glbPath_output,objPath_output, info = _gen_shape_on_gpu(
+        image, 
+        steps,
+        guidance_scale, 
+        seed,
+        octree_resolution,
+        num_chunks, 
+        target_face_num,
+        randomize_seed
+    )
+    # 如果出错，抛出异常
+    if info["status"] == "failed":
+        raise gr.Error(info["error"])
+    return html_export_mesh, file_export, glbPath_output, objPath_output
     
 def get_example_img_list():
     print('Loading example img list ...')