tianleliphoebe commited on
Commit
85dbc3c
1 Parent(s): 52026b9

add NSFW guard

Browse files
Files changed (1) hide show
  1. model/model_manager.py +10 -4
model/model_manager.py CHANGED
@@ -52,8 +52,11 @@ class ModelManager:
52
  return result
53
 
54
  def generate_image_ig_api(self, prompt, model_name):
55
- pipe = self.load_model_pipe(model_name)
56
- result = pipe(prompt=prompt)
 
 
 
57
  return result
58
 
59
  def generate_image_ig_museum(self, model_name):
@@ -192,8 +195,11 @@ class ModelManager:
192
  return result
193
 
194
  def generate_video_vg_api(self, prompt, model_name):
195
- pipe = self.load_model_pipe(model_name)
196
- result = pipe(prompt=prompt)
 
 
 
197
  return result
198
 
199
  def generate_video_vg_museum(self, model_name):
 
52
  return result
53
 
54
  def generate_image_ig_api(self, prompt, model_name):
55
+ if self.NSFW_filter(prompt) == 'safe':
56
+ pipe = self.load_model_pipe(model_name)
57
+ result = pipe(prompt=prompt)
58
+ else:
59
+ result = ''
60
  return result
61
 
62
  def generate_image_ig_museum(self, model_name):
 
195
  return result
196
 
197
  def generate_video_vg_api(self, prompt, model_name):
198
+ if self.NSFW_filter(prompt) == 'safe':
199
+ pipe = self.load_model_pipe(model_name)
200
+ result = pipe(prompt=prompt)
201
+ else:
202
+ result = ''
203
  return result
204
 
205
  def generate_video_vg_museum(self, model_name):