Update app.py

app.py

@@ -76,7 +76,8 @@ def encode_prompt(
     prompt: str,
     enhance_prompt: bool = True,
     input_image = None,
-    seed: int = 42
+    seed: int = 42,
+    negative_prompt: str = ""
 ):
     """
     Encode a text prompt using Gemma text encoder.
@@ -86,6 +87,7 @@ def encode_prompt(
         enhance_prompt: Whether to use AI to enhance the prompt
         input_image: Optional image for image-to-video enhancement
         seed: Random seed for prompt enhancement
+        negative_prompt: Optional negative prompt for CFG (two-stage pipeline)
 
     Returns:
         tuple: (file_path, enhanced_prompt_text, status_message)
@@ -120,21 +122,35 @@ def encode_prompt(
                     seed=seed
                 )
 
-        # Encode the prompt using the pre-loaded text encoder
+        # Encode the positive prompt using the pre-loaded text encoder
         video_context, audio_context = encode_text_simple(text_encoder, final_prompt)
 
+        # Encode negative prompt if provided
+        video_context_negative = None
+        audio_context_negative = None
+        if negative_prompt:
+            video_context_negative, audio_context_negative = encode_text_simple(text_encoder, negative_prompt)
+
         # Save embeddings to file
         output_dir = Path("embeddings")
         output_dir.mkdir(exist_ok=True)
         output_path = output_dir / f"embedding_{int(time.time())}.pt"
 
-        # Save both video and audio contexts
-        torch.save({
+        # Save embeddings (with negative contexts if provided)
+        embedding_data = {
             'video_context': video_context.cpu(),
             'audio_context': audio_context.cpu(),
             'prompt': final_prompt,
             'original_prompt': prompt if enhance_prompt else final_prompt,
-        }, output_path)
+        }
+
+        # Add negative contexts if they were encoded
+        if video_context_negative is not None:
+            embedding_data['video_context_negative'] = video_context_negative.cpu()
+            embedding_data['audio_context_negative'] = audio_context_negative.cpu()
+            embedding_data['negative_prompt'] = negative_prompt
+
+        torch.save(embedding_data, output_path)
 
         # Get memory stats
         elapsed_time = time.time() - start_time
@@ -171,6 +187,13 @@ with gr.Blocks(title="LTX-2 Gemma Text Encoder") as demo:
                 value="An astronaut hatches from a fragile egg on the surface of the Moon"
             )
 
+            negative_prompt_input = gr.Textbox(
+                label="Negative Prompt (Optional)",
+                placeholder="Enter negative prompt for CFG (used by two-stage pipeline)...",
+                lines=2,
+                value=""
+            )
+
             enhance_checkbox = gr.Checkbox(
                 label="Enhance Prompt",
                 value=True,
@@ -204,7 +227,7 @@ with gr.Blocks(title="LTX-2 Gemma Text Encoder") as demo:
 
     encode_btn.click(
         fn=encode_prompt,
-        inputs=[prompt_input, enhance_checkbox, input_image, enhancement_seed],
+        inputs=[prompt_input, enhance_checkbox, input_image, enhancement_seed, negative_prompt_input],
         outputs=[embedding_file, enhanced_prompt_output, status_output]
     )