Update app.py

app.py

@@ -31,10 +31,6 @@ from diffusers import (
 from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
 from diffusers.utils.export_utils import export_to_video
 
-# Safety Checker Imports
-from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
-from transformers import AutoImageProcessor
-
 from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig, Int8WeightOnlyConfig
 import aoti
 
@@ -46,18 +42,6 @@ IS_ZERO_GPU = bool(os.getenv("SPACES_ZERO_GPU"))
 #     print("Loading...")
 #     subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
 
-# --- SAFETY CHECKER INITIALIZATION ---
-safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to("cuda")
-feature_extractor = AutoImageProcessor.from_pretrained("CompVis/stable-diffusion-safety-checker")
-
-def check_nsfw(frame_np):
-    """Checks a single numpy frame (float32, 0-1) for NSFW content."""
-    # Convert back to uint8 PIL for the processor
-    img = Image.fromarray((frame_np * 255).astype(np.uint8))
-    inputs = feature_extractor(img, return_tensors="pt").to("cuda")
-    # Safety checker expects float16 on GPU
-    _, has_nsfw_concept = safety_checker(images=inputs.pixel_values, clip_input=inputs.pixel_values.to(torch.float16))
-    return has_nsfw_concept[0]
 
 # --- FRAME EXTRACTION JS & LOGIC ---
 
@@ -373,7 +357,6 @@ def get_inference_duration(
     quality,
     duration_seconds,
     safe_mode,
-    enable_safety_checker,
     progress
 ):
     BASE_FRAMES_HEIGHT_WIDTH = 81 * 832 * 624
@@ -417,7 +400,6 @@ def run_inference(
     quality,
     duration_seconds,
     safe_mode,
-    enable_safety_checker,
     progress=gr.Progress(track_tqdm=True),
 ):
     scheduler_class = SCHEDULER_MAP.get(scheduler_name)
@@ -452,18 +434,6 @@ def run_inference(
     raw_frames_np = result.frames[0]  # Returns (T, H, W, C) float32
     pipe.scheduler = original_scheduler
 
-    # --- SAFETY CHECKER LOGIC ---
-    is_nsfw = False
-    if enable_safety_checker:
-        # Implements an automated review process to check the results before making them publicly available.
-        # If the user did not supply a 'last_image', the final freely-generated frame is verified by 
-        # the safety checker to ensure no unrequested explicit content was created.
-        if processed_last_image is None:
-            is_nsfw = check_nsfw(raw_frames_np[-1])
-            
-    if is_nsfw:
-        return None, task_name, True
-
     frame_factor = frame_multiplier // FIXED_FPS
     if frame_factor > 1:
         start = time.time()
@@ -484,7 +454,7 @@ def run_inference(
         export_to_video(final_frames, video_path, fps=final_fps, quality=quality)
         pbar.update(1)
 
-    return video_path, task_name, False
+    return video_path, task_name
 
 
 def generate_video(
@@ -504,54 +474,8 @@ def generate_video(
     frame_multiplier=16,
     video_component=True,
     safe_mode=False,
-    enable_safety_checker=True,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """
-    Generate a video from an input image using the Wan 2.2 14B I2V model with Lightning LoRA.
-    This function takes an input image and generates a video animation based on the provided
-    prompt and parameters. It uses an FP8 qunatized Wan 2.2 14B Image-to-Video model in with Lightning LoRA
-    for fast generation in 4-8 steps.
-    Args:
-        input_image (PIL.Image): The input image to animate. Will be resized to target dimensions.
-        last_image (PIL.Image, optional): The optional last image for the video.
-        prompt (str): Text prompt describing the desired animation or motion.
-        steps (int, optional): Number of inference steps. More steps = higher quality but slower.
-            Defaults to 4. Range: 1-30.
-        negative_prompt (str, optional): Negative prompt to avoid unwanted elements.
-            Defaults to default_negative_prompt (contains unwanted visual artifacts).
-        duration_seconds (float, optional): Duration of the generated video in seconds.
-            Defaults to 2. Clamped between MIN_FRAMES_MODEL/FIXED_FPS and MAX_FRAMES_MODEL/FIXED_FPS.
-        guidance_scale (float, optional): Controls adherence to the prompt. Higher values = more adherence.
-            Defaults to 1.0. Range: 0.0-20.0.
-        guidance_scale_2 (float, optional): Controls adherence to the prompt. Higher values = more adherence.
-            Defaults to 1.0. Range: 0.0-20.0.
-        seed (int, optional): Random seed for reproducible results. Defaults to 42.
-            Range: 0 to MAX_SEED (2147483647).
-        randomize_seed (bool, optional): Whether to use a random seed instead of the provided seed.
-            Defaults to False.
-        quality (float, optional): Video output quality. Default is 5. Uses variable bit rate.
-            Highest quality is 10, lowest is 1.
-        scheduler (str, optional): The name of the scheduler to use for inference. Defaults to "UniPCMultistep".
-        flow_shift (float, optional): The flow shift value for compatible schedulers. Defaults to 6.0.
-        frame_multiplier (int, optional): The int value for fps enhancer
-        video_component(bool, optional): Show video player in output.
-            Defaults to True.
-        enable_safety_checker(bool, optional): Enable NSFW filter.
-        progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
-    Returns:
-        tuple: A tuple containing:
-            - video_path (str): Path for the video component.
-            - video_path (str): Path for the file download component. Attempt to avoid reconversion in video component.
-            - current_seed (int): The seed used for generation.
-    Raises:
-        gr.Error: If input_image is None (no image uploaded).
-    Note:
-        - Frame count is calculated as duration_seconds * FIXED_FPS (24)
-        - Output dimensions are adjusted to be multiples of MOD_VALUE (32)
-        - The function uses GPU acceleration via the @spaces.GPU decorator
-        - Generation time varies based on steps and duration (see get_duration function)
-    """
     
     if input_image is None:
         raise gr.Error("Please upload an input image.")
@@ -564,7 +488,7 @@ def generate_video(
     if last_image:
         processed_last_image = resize_and_crop_to_match(last_image, resized_image)
 
-    video_path, task_n, is_nsfw = run_inference(
+    video_path, task_n = run_inference(
         resized_image,
         processed_last_image,
         prompt,
@@ -580,13 +504,9 @@ def generate_video(
         quality,
         duration_seconds,
         safe_mode,
-        enable_safety_checker,
         progress,
     )
 
-    if is_nsfw:
-        raise gr.Error("Generation blocked by guardrails: The resulting video may contain explicit content.")
-
     print(f"GPU complete: {task_n}")
 
     return (video_path if video_component else None), video_path, current_seed
@@ -643,7 +563,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CSS, delete_cache=(3600, 3700)) as de
                     info="Select a custom scheduler."
                 )
                 flow_shift_slider = gr.Slider(minimum=0.5, maximum=15.0, step=0.1, value=3.0, label="Flow Shift")
-                safety_checker_input = gr.Checkbox(label="Enable Safety Filter", value=True, info="Prevents unrequested sensitive or explicit content.")
                 play_result_video = gr.Checkbox(label="Display result", value=True, interactive=True)
 
             generate_button = gr.Button("Generate Video", variant="primary")
@@ -662,37 +581,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CSS, delete_cache=(3600, 3700)) as de
 
     ui_inputs =[
         input_image_component, last_image_component, prompt_input, steps_slider,
-        negative_prompt_input, duration_seconds_input,
-        guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox,
-        quality_slider, scheduler_dropdown, flow_shift_slider, frame_multi,
-        play_result_video, safe_mode_checkbox, safety_checker_input
-    ]
-    
-    generate_button.click(
-        fn=generate_video, 
-        inputs=ui_inputs, 
-        outputs=[video_output, file_output, seed_input]
-    )
-    
-    # --- Frame Grabbing Events ---
-    # 1. Click button -> JS runs -> puts time in hidden number box
-    grab_frame_btn.click(
-        fn=None,
-        inputs=None,
-        outputs=[timestamp_box],
-        js=get_timestamp_js
-    )
-    
-    # 2. Hidden number box changes -> Python runs -> puts frame in Input Image
-    timestamp_box.change(
-        fn=extract_frame,
-        inputs=[video_output, timestamp_box],
-        outputs=[input_image_component]
-    )
-
-if __name__ == "__main__":
-    demo.queue().launch(
-        mcp_server=True,
-        ssr_mode=False,
-        show_error=True,
-    )
+        neg