Update app.py

app.py

@@ -33,7 +33,7 @@ from PIL import Image
 from torchvision.utils import make_grid, save_image
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from app import safety_check
+#from app import safety_check
 from app.sana_pipeline import SanaPipeline
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -205,12 +205,12 @@ if torch.cuda.is_available():
     pipe.register_progress_bar(gr.Progress())
 
     # safety checker
-    safety_checker_tokenizer = AutoTokenizer.from_pretrained(args.shield_model_path)
-    safety_checker_model = AutoModelForCausalLM.from_pretrained(
-        args.shield_model_path,
-        device_map="auto",
-        torch_dtype=torch.bfloat16,
-    ).to(device)
+    #safety_checker_tokenizer = AutoTokenizer.from_pretrained(args.shield_model_path)
+    #safety_checker_model = AutoModelForCausalLM.from_pretrained(
+    #    args.shield_model_path,
+    #    device_map="auto",
+    #    torch_dtype=torch.bfloat16,
+    #).to(device)
 
 
 def save_image_sana(img, seed="", save_img=False):
@@ -254,8 +254,8 @@ def generate(
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
     print(f"PORT: {DEMO_PORT}, model_path: {model_path}")
-    if safety_check.is_dangerous(safety_checker_tokenizer, safety_checker_model, prompt, threshold=0.2):
-        prompt = "A red heart."
+    #if safety_check.is_dangerous(safety_checker_tokenizer, safety_checker_model, prompt, threshold=0.2):
+    #    prompt = "A red heart."
 
     print(prompt)