image to 3d

app.py

@@ -12,8 +12,6 @@ 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')
@@ -40,11 +38,6 @@ 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:
@@ -138,22 +131,10 @@ 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_on_gpu(
+def gen_shape(
     image=None,
     steps=50,
     guidance_scale=7.5,
@@ -171,28 +152,14 @@ def _gen_shape_on_gpu(
 
 
     if image is None:
-        error_info = {
-            "error": "Please provide either a caption or an image.",
-            "status": "failed",
-        }
-        return None,None,None,None,error_info
+        raise gr.Error("Please provide either a caption or an image.")
 
-    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(rgbImage)
+    image = rmbg_worker(image.convert('RGB'))
 
     # 生成模型
     generator = torch.Generator()
@@ -221,11 +188,7 @@ def _gen_shape_on_gpu(
         torch.cuda.empty_cache()
     
     if path is None:
-        error_info = {
-            "error": "'Please generate a mesh first.'",
-            "status": "failed",
-        }
-        return None,None,None,None,error_info
+        raise gr.Error('Please generate a mesh first.')
 
     # 简化模型
     print(f'exporting {path}')
@@ -258,37 +221,9 @@ def _gen_shape_on_gpu(
 
 
     progress(1,desc="Complete")
-    info = {
-        "status": "success"
-    }
-    return model_viewer_html, gr.update(value=sourceObjPath, interactive=True), glbPath, objPath, info
+    return model_viewer_html, gr.update(value=sourceObjPath, interactive=True), glbPath, objPath
 
 
-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 ...')

requirements.txt

@@ -18,7 +18,7 @@ tqdm
 
 # Mesh Processing
 trimesh
-pymeshlab==2023.12.post3
+pymeshlab
 pygltflib
 xatlas
 #kornia