Update app.py

app.py

@@ -67,44 +67,40 @@ def translate_language(text_prompts):
         return {language_tips_text:gr.update(visible=True, value=tips_text), translated_language:text_prompts, trigger_component:  gr.update(value=count, visible=False)}
 
 
+word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)
+word_list = word_list_dataset["train"]['text']
+
+
 def get_result(text_prompts, style_indx):
-  #results = text_to_img(text_prompts, style_indx, fn_index=3)
-  print(text_prompts)
-  #results = text_to_img(text_prompts, text_prompts, fn_index=3)
-  print(style_indx)
-  try:
     style = style_list[style_indx]
-      
-    #results = text_to_img(text_prompts, style_indx, fn_index=3)
-      
-    results = model.generate_image(
-        text_prompts=text_prompts, style=style, visualization=False, topk=1)
-  except Exception as e:
-    error_text = str(e)
-    return {video_result:None, status_text:error_text}
-  
-  #print("Ernie Vilg Output: " + str(results[:1]))
-  print("Ernie Vilg Output test: " + str(results))
+    prompt = "," + style
+    for filter in word_list:
+        if re.search(rf"\b{filter}\b", text_prompts):
+            raise gr.Error("Unsafe content found. Please try again with different prompts.")
+        
+    model_id = "runwayml/stable-diffusion-v1-5"
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+    pipe = pipe.to("cuda")
     
-  #image_output = results[:1]
-  image_output = results[:1]
-  
-  print("file name: " + image_output[0].filename)
-
-  # Encode your PIL Image as a JPEG without writing to disk
-  imagefile = "imageoutput.png"
-  #img_np = np.array(image_output[0])
-  #img_nparray= cv2.cvtColor(img_np, cv2.COLOR_BGR2RGBA)
-  #img_blue_correction = Image.fromarray(img_nparray)
-  #img_blue_correction.save(imagefile, img_blue_correction.format)
-  image_output[0].save(imagefile, image_output[0].format)
-
-  interrogate_prompt = img_to_text(imagefile, fn_index=1)[0]
-  print(interrogate_prompt)
-  spec_image, music_output = get_music(interrogate_prompt + ", " + style_list_EN[style_indx])
-  
-  video_merged = merge_video(music_output, image_output)
-  return {spec_result:spec_image, video_result:video_merged, status_text:'Success'}
+    prompt = text_prompts
+    image_output = pipe(prompt).images[0]
+    
+    print("file name: " + image_output.filename)
+    
+    # Encode your PIL Image as a JPEG without writing to disk
+    imagefile = "imageoutput.png"
+    #img_np = np.array(image_output[0])
+    #img_nparray= cv2.cvtColor(img_np, cv2.COLOR_BGR2RGBA)
+    #img_blue_correction = Image.fromarray(img_nparray)
+    #img_blue_correction.save(imagefile, img_blue_correction.format)
+    image_output[0].save(imagefile, image_output[0].format)
+    
+    interrogate_prompt = img_to_text(imagefile, fn_index=1)[0]
+    print(interrogate_prompt)
+    spec_image, music_output = get_music(interrogate_prompt + ", " + style_list_EN[style_indx])
+    
+    video_merged = merge_video(music_output, image_output)
+    return {spec_result:spec_image, video_result:video_merged, status_text:'Success'}
 
 def get_music(prompt):
   spec = pipe(prompt).images[0]