| from typing import List |
| from diffusers.modular_pipelines import ( |
| PipelineState, |
| ModularPipelineBlocks, |
| InputParam, |
| OutputParam, |
| ) |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| 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 QwenPromptExpander(ModularPipelineBlocks): |
| def __init__(self, model_id="Qwen/Qwen2.5-3B-Instruct", system_prompt=SYSTEM_PROMPT): |
| super().__init__() |
|
|
| self.system_prompt = system_prompt |
| self.model = AutoModelForCausalLM.from_pretrained( |
| model_id, torch_dtype="auto" |
| ).to("cuda") |
| self.tokenizer = AutoTokenizer.from_pretrained(model_id) |
|
|
| @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_inputs(self) -> List[InputParam]: |
| return [] |
|
|
| @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) |
|
|
| old_prompt = block_state.prompt |
| print(f"Actual prompt: {old_prompt}") |
| |
| messages = [ |
| {"role": "system", "content": self.system_prompt}, |
| {"role": "user", "content": old_prompt} |
| ] |
| text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device) |
| generated_ids = self.model.generate(**model_inputs,max_new_tokens=512) |
| generated_ids = [ |
| output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids) |
| ] |
|
|
| block_state.prompt = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] |
| block_state.old_prompt = old_prompt |
| print(f"{block_state.prompt=}") |
| self.set_block_state(state, block_state) |
|
|
| return components, state |
