commit

app.py

@@ -124,17 +124,24 @@ function custom(){
 
 # torch pipes
 
-image_pipe = FluxPipeline.from_pretrained(base, torch_dtype=torch.bfloat16).to(device)
+def disabled_safety_checker(images, clip_input):
+    if len(images.shape)==4:
+        num_images = images.shape[0]
+        return images, [False]*num_images
+    else:
+        return images, False
 
+image_pipe = FluxPipeline.from_pretrained(base, torch_dtype=torch.bfloat16).to(device)
 image_pipe.enable_model_cpu_offload()
+image_pipe.safety_checker = disabled_safety_checker
 
 video_pipe = CogVideoXImageToVideoPipeline.from_pretrained(
     "THUDM/CogVideoX-5b-I2V",
     torch_dtype=torch.bfloat16
 ).to(device)
-
 video_pipe.vae.enable_tiling()
 video_pipe.vae.enable_slicing()
+video_pipe.safety_checker = disabled_safety_checker
 
 # functionality
 
@@ -155,7 +162,6 @@ def pipe_generate(img,p1,p2,time,title):
             guidance_scale=img_accu,
             num_images_per_prompt=1,
             num_inference_steps=image_steps,
-            safety_checker=None,
             max_sequence_length=seq,
             generator=torch.Generator(device).manual_seed(int(str(random.random()).split(".")[1]))
         ).images[0]
@@ -178,7 +184,6 @@ def pipe_generate(img,p1,p2,time,title):
 
     return video_pipe(
         prompt=p1,
-        safety_checker=None,
         negative_prompt=p2.replace("textual content, ",""),
         image=img,
         num_inference_steps=video_steps,