| """Helpers shared by the Kimodo blocks.""" |
|
|
| from dataclasses import replace |
|
|
| import torch |
| from diffusers.modular_pipelines import ModularPipelineBlocks |
|
|
|
|
| class ResolvableBlock(ModularPipelineBlocks): |
| """Leaf block whose generic component specs can carry the model repository's resolved classes and location. |
| |
| `KimodoBlocks(components_repo=...)` fills `component_overrides`, so a sub-block popped out of the pipeline still |
| loads the right custom classes from the right repository. |
| """ |
|
|
| component_overrides = {} |
|
|
| def resolve(self, specs): |
| return [ |
| replace(spec, **self.component_overrides[spec.name]) if spec.name in self.component_overrides else spec |
| for spec in specs |
| ] |
|
|
|
|
| class Cancelled(Exception): |
| pass |
|
|
|
|
| def checkpoint(cancel): |
| if cancel is not None and cancel.is_set(): |
| raise Cancelled("Generation cancelled") |
|
|
|
|
| @torch.no_grad() |
| def encode_prompt(text_encoder, tokenizer, prompt, device): |
| """Encode one prompt into LLM2Vec features shaped (1, 1, hidden).""" |
| from .prompt_utils import tokenize_prompt |
| from .sanitize import sanitize_texts |
|
|
| inputs = tokenize_prompt(tokenizer, sanitize_texts([prompt])[0]) |
| features = text_encoder(**{name: value.to(text_encoder.device) for name, value in inputs.items()}) |
| return features.to(device) |
|
|
