|
|
| 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. " |
| "Please don't return any prefix or suffix tokens, just the expanded user description." |
| ) |
|
|
| class GeminiPromptExpander(ModularPipelineBlocks): |
| model_name = "flux" |
| |
| 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", |
| metadata={"mellon":"textbox"} |
| ) |
| ] |
| |
| @property |
| def intermediate_inputs(self) -> List[InputParam]: |
| return [] |
|
|
| @property |
| def intermediate_outputs(self) -> List[OutputParam]: |
| return [ |
| OutputParam( |
| "prompt", |
| type_hint=str, |
| description="Expanded prompt by the LLM", |
| metadata={"mellon":"text"} |
| ), |
| 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) |
|
|
| 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 |
|
|
