vaibhaviiii28 commited on
Commit
5fe10ca
Β·
verified Β·
1 Parent(s): 77e84bd

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -4
app.py CHANGED
@@ -2,6 +2,7 @@ import os
2
  from fastapi import FastAPI, File, UploadFile
3
  from pydantic import BaseModel
4
  from transformers import pipeline
 
5
  from PIL import Image
6
  import joblib
7
  import re
@@ -13,11 +14,9 @@ import uvicorn
13
  os.environ["TRANSFORMERS_CACHE"] = "/tmp"
14
  os.environ["HF_HOME"] = "/tmp"
15
 
16
- # βœ… Initialize FastAPI
17
- app = FastAPI()
18
-
19
- # βœ… Load NSFW Image Classification Model
20
  try:
 
21
  pipe = pipeline("image-classification", model="LukeJacob2023/nsfw-image-detector", cache_dir="/tmp")
22
  print("βœ… NSFW Model Loaded Successfully!")
23
  except Exception as e:
@@ -33,6 +32,9 @@ except Exception as e:
33
  print(f"❌ Error Loading Toxic Text Model: {e}")
34
  exit(1)
35
 
 
 
 
36
  # πŸ“Œ Text Input Model
37
  class TextInput(BaseModel):
38
  text: str
 
2
  from fastapi import FastAPI, File, UploadFile
3
  from pydantic import BaseModel
4
  from transformers import pipeline
5
+ from huggingface_hub import hf_hub_download
6
  from PIL import Image
7
  import joblib
8
  import re
 
14
  os.environ["TRANSFORMERS_CACHE"] = "/tmp"
15
  os.environ["HF_HOME"] = "/tmp"
16
 
17
+ # βœ… Manually Download the NSFW Model to `/tmp`
 
 
 
18
  try:
19
+ model_path = hf_hub_download(repo_id="LukeJacob2023/nsfw-image-detector", filename="pytorch_model.bin", cache_dir="/tmp")
20
  pipe = pipeline("image-classification", model="LukeJacob2023/nsfw-image-detector", cache_dir="/tmp")
21
  print("βœ… NSFW Model Loaded Successfully!")
22
  except Exception as e:
 
32
  print(f"❌ Error Loading Toxic Text Model: {e}")
33
  exit(1)
34
 
35
+ # βœ… Initialize FastAPI
36
+ app = FastAPI()
37
+
38
  # πŸ“Œ Text Input Model
39
  class TextInput(BaseModel):
40
  text: str