Create prompt_expander.py

prompt_expander.py

@@ -0,0 +1,67 @@
+
+from typing import List
+from diffusers.modular_pipelines import (
+    PipelineState,
+    ModularPipelineBlocks,
+    InputParam,
+    OutputParam,
+)
+import google.generativeai as genai
+import os
+
+SYSTEM_PROMPT = (
+    "You are an expert image generation assistant. "
+    "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."
+)
+
+class GeminiPromptExpander(ModularPipelineBlocks):
+    def __init__(self, model_id="gemini-2.5-flash-lite", system_prompt=SYSTEM_PROMPT):
+        super().__init__()
+        api_key = os.getenv("GOOGLE_API_KEY")
+        if api_key is None:
+            raise ValueError("Must provide an API key for Gemini through the `GOOGLE_API_KEY` env variable.")
+        genai.configure(api_key=api_key)
+        self.model = genai.GenerativeModel(model_name=model_id, system_instruction=system_prompt)
+
+    @property
+    def expected_components(self):
+        return []
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam(
+                "prompt",
+                type_hint=str,
+                required=True,
+                description="Prompt to use",
+            )
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam(
+                "prompt",
+                type_hint=str,
+                description="Expanded prompt by the LLM",
+            ),
+            OutputParam(
+                "old_prompt",
+                type_hint=str,
+                description="Old prompt provided by the user",
+            )
+        ]
+
+
+    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
+        self.set_block_state(state, block_state)
+
+        return components, state