add password for generation

app.py

@@ -1,4 +1,5 @@
 import os
+import hmac
 import spaces
 import torch
 from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
@@ -28,6 +29,7 @@ MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
 DATASET_REPO_ID = os.environ.get("DATASET_REPO_ID")
 PRIVATE_MODEL_KEY = os.environ.get("PRIVATE_MODEL_KEY")
 hf_api = HfApi(token=key) if key and DATASET_REPO_ID else None
+GENERATION_SECRET=os.environ.get("GENERATION_SECRET")
 
 MAX_DIM = 832
 MIN_DIM = 480
@@ -247,6 +249,7 @@ def get_duration(
     guidance_scale_2,
     seed,
     randomize_seed,
+    access_granted,
     request,
     progress,
 ):
@@ -325,6 +328,39 @@ def export_to_video_h264(frames, output_path: str, fps: int = FIXED_FPS) -> None
             writer.append_data(frame)
 
 
+def reset_verification():
+    return (
+        gr.update(value="Verify Passcode to Enable Generation", interactive=False),
+        "Enter the passcode and click Verify to unlock generation.",
+        False,
+    )
+
+
+def verify_passcode(passcode: str):
+    secret = GENERATION_SECRET or ""
+    candidate = (passcode or "").strip()
+
+    if not secret:
+        return (
+            gr.update(value="Verify Passcode to Enable Generation", interactive=False),
+            "Generation is unavailable because `GENERATION_SECRET` is not configured.",
+            False,
+        )
+
+    if hmac.compare_digest(candidate, secret):
+        return (
+            gr.update(value="Generate Video", interactive=True),
+            "Passcode verified. Generation unlocked.",
+            True,
+        )
+
+    return (
+        gr.update(value="Verify Passcode to Enable Generation", interactive=False),
+        "Incorrect passcode.",
+        False,
+    )
+
+
 @spaces.GPU(duration=get_duration)
 def generate_video(
     input_image,
@@ -336,6 +372,7 @@ def generate_video(
     guidance_scale_2 = 1,    
     seed = 7777777,
     randomize_seed = False,
+    access_granted = False,
     request: gr.Request = None,
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -379,6 +416,9 @@ def generate_video(
         - The function uses GPU acceleration via the @spaces.GPU decorator
         - Generation time varies based on steps and duration (see get_duration function)
     """
+    if not access_granted:
+        raise gr.Error("Please verify the passcode before generating.")
+
     if input_image is None:
         raise gr.Error("Please upload an input image.")
     
@@ -438,6 +478,14 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             input_image_component = gr.Image(type="filepath", label="Input Image")
+            verification_state = gr.State(False)
+            passcode_input = gr.Textbox(
+                label="Passcode",
+                type="password",
+                placeholder="Enter passcode to unlock generation",
+            )
+            verify_button = gr.Button("Verify", variant="secondary")
+            verification_status = gr.Markdown("Enter the passcode and click Verify to unlock generation.")
             prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v, lines=3)
             duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=3.5, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
             
@@ -445,20 +493,30 @@ with gr.Blocks() as demo:
                 negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                 seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=7777777, interactive=True)
                 randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
-                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps") 
-                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage")
-                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2 - low noise stage")
+                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps", interactive=True) 
+                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage", interactive=True)
+                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2 - low noise stage", interactive=True)
 
-            generate_button = gr.Button("Generation Available on Patreon", variant="primary",interactive=False)
+            generate_button = gr.Button("Verify Passcode to Enable Generation", variant="primary", interactive=False)
         with gr.Column():
             video_output = gr.Video(label="Generated Video", autoplay=True, sources=["upload"], interactive=False, elem_id="generated-video")
     
     ui_inputs = [
         input_image_component, prompt_input, steps_slider,
         negative_prompt_input, duration_seconds_input,
-        guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox
+        guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox,
+        verification_state,
     ]
-    # generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
+    passcode_input.change(
+        fn=reset_verification,
+        outputs=[generate_button, verification_status, verification_state],
+    )
+    verify_button.click(
+        fn=verify_passcode,
+        inputs=[passcode_input],
+        outputs=[generate_button, verification_status, verification_state],
+    )
+    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
 
     gr.Examples(
         examples=[ 
@@ -498,10 +556,7 @@ with gr.Blocks() as demo:
                 6,
             ],
         ],
-        inputs=[input_image_component, prompt_input, steps_slider], 
-        outputs=[video_output, seed_input], 
-        fn=generate_video, 
-        cache_examples="lazy"
+        inputs=[input_image_component, prompt_input, steps_slider],
     )
     gr.Markdown("*Content Warning:* This research demonstration is capable of generating content that may not be suitable for all audiences. By using this application, you agree to assume full responsibility for your compliance with the platform's terms of service and any applicable legal and regulatory guidelines.")