sayakpaul HF Staff commited on
Commit
d84a752
·
verified ·
1 Parent(s): bbf66ad

Update prompt_expander.py

Browse files
Files changed (1) hide show
  1. prompt_expander.py +4 -2
prompt_expander.py CHANGED
@@ -19,6 +19,8 @@ SYSTEM_PROMPT = (
19
  )
20
 
21
  class GeminiPromptExpander(ModularPipelineBlocks):
 
 
22
  def __init__(self, model_id="gemini-2.5-flash-lite", system_prompt=SYSTEM_PROMPT):
23
  super().__init__()
24
  api_key = os.getenv("GOOGLE_API_KEY")
@@ -66,10 +68,10 @@ class GeminiPromptExpander(ModularPipelineBlocks):
66
  block_state = self.get_block_state(state)
67
 
68
  old_prompt = block_state.prompt
69
- # print(f"Actual prompt: {old_prompt}")
70
  block_state.prompt = self.model.generate_content(old_prompt).text
71
  block_state.old_prompt = old_prompt
72
- # print(f"{block_state.prompt=}")
73
  self.set_block_state(state, block_state)
74
 
75
  return components, state
 
19
  )
20
 
21
  class GeminiPromptExpander(ModularPipelineBlocks):
22
+ model_name = "flux"
23
+
24
  def __init__(self, model_id="gemini-2.5-flash-lite", system_prompt=SYSTEM_PROMPT):
25
  super().__init__()
26
  api_key = os.getenv("GOOGLE_API_KEY")
 
68
  block_state = self.get_block_state(state)
69
 
70
  old_prompt = block_state.prompt
71
+ print(f"Actual prompt: {old_prompt}")
72
  block_state.prompt = self.model.generate_content(old_prompt).text
73
  block_state.old_prompt = old_prompt
74
+ print(f"{block_state.prompt=}")
75
  self.set_block_state(state, block_state)
76
 
77
  return components, state