Update prompt_expander.py

prompt_expander.py

@@ -14,7 +14,8 @@ SYSTEM_PROMPT = (
     "Take the user's short description and expand it into a vivid, detailed, and clear image generation prompt. "
     "Ensure rich colors, depth, realistic lighting, and an imaginative composition. "
     "Avoid vague terms — be specific about style, perspective, and mood. "
-    "Try to keep the output under 512 tokens."
+    "Try to keep the output under 512 tokens. "
+    "Please don't return any prefix or suffix tokens, just the expanded user description."
 )
 
 class GeminiPromptExpander(ModularPipelineBlocks):
@@ -64,8 +65,11 @@ class GeminiPromptExpander(ModularPipelineBlocks):
     def __call__(self, components, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
 
-        block_state.old_prompt = block_state.prompt
-        block_state.prompt = self.model.generate_content(block_state.old_prompt).text
+        old_prompt = block_state.prompt
+        # print(f"Actual prompt: {old_prompt}")
+        block_state.prompt = self.model.generate_content(old_prompt).text
+        block_state.old_prompt = old_prompt
+        # print(f"{block_state.prompt=}")
         self.set_block_state(state, block_state)
 
         return components, state