Update Athena runtimes and sync live extensions
This commit is contained in:
Mikei386
@@ -0,0 +1,354 @@
|
||||
"""Local, catalog-aware prompt enhancement handler."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import uuid
|
||||
from threading import RLock
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from _routes._errors import HTTPError
|
||||
from api_types import EnhancePromptRequest, EnhancePromptResponse, IcLoraCatalogItem, LoraCatalogItem
|
||||
from handlers.base import StateHandlerBase
|
||||
from handlers.generation_handler import GenerationHandler
|
||||
from handlers.pipelines_handler import PipelinesHandler
|
||||
from handlers.text_handler import TextHandler
|
||||
from server_utils.media_validation import normalize_optional_path, validate_image_file
|
||||
from services.gemini_text_client import resolve_gemini_model
|
||||
from services.interfaces import PromptEnhancerPipeline
|
||||
from services.lora_catalog import LoraCatalogProvider
|
||||
from services.prompt_enhancement import (
|
||||
build_audio_visual_caption_system_prompt,
|
||||
build_conditioning_system_prompt,
|
||||
build_ic_lora_enhancement_system_prompt,
|
||||
build_image_edit_system_prompt,
|
||||
build_image_generation_system_prompt,
|
||||
build_keyframe_enhancement_system_prompt,
|
||||
build_lora_enhancement_system_prompt,
|
||||
build_template_fill_system_prompt,
|
||||
enforce_trigger_placements,
|
||||
fill_prompt_template,
|
||||
parse_template_fill_response,
|
||||
)
|
||||
from services.prompt_enhancement.i2v_frames import KeyframeStill
|
||||
from services.prompt_enhancer_pipeline.gemini_prompt_enhancer_pipeline import GeminiPromptEnhancerPipeline
|
||||
from services.services_utils import get_device_type
|
||||
from state.app_state_types import AppState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from runtime_config.runtime_config import RuntimeConfig
|
||||
|
||||
# Deliberately independent of StateHandlerBase._resolve_seed(): that helper honors the app's
|
||||
# reproducibility seed lock (and a fixed constant in dev mode), which would make every enhance
|
||||
# call — including a redo — produce the exact same output. Enhancement is a quick, exploratory
|
||||
# action where a fresh draw each call is the whole point.
|
||||
_MAX_ENHANCE_SEED = 2147483647
|
||||
|
||||
|
||||
class PromptEnhancementHandler(StateHandlerBase):
|
||||
def __init__(
|
||||
self,
|
||||
state: AppState,
|
||||
lock: RLock,
|
||||
generation_handler: GenerationHandler,
|
||||
pipelines_handler: PipelinesHandler,
|
||||
text_handler: TextHandler,
|
||||
lora_catalog_provider: LoraCatalogProvider,
|
||||
prompt_enhancer_pipeline_class: type[PromptEnhancerPipeline],
|
||||
gemini_pipeline: GeminiPromptEnhancerPipeline,
|
||||
config: RuntimeConfig,
|
||||
) -> None:
|
||||
super().__init__(state, lock, config)
|
||||
self._generation = generation_handler
|
||||
self._pipelines = pipelines_handler
|
||||
self._text_handler = text_handler
|
||||
self._lora_catalog_provider = lora_catalog_provider
|
||||
self._prompt_enhancer_pipeline_class = prompt_enhancer_pipeline_class
|
||||
self._gemini_pipeline = gemini_pipeline
|
||||
|
||||
def _random_seed(self) -> int:
|
||||
return random.randint(0, _MAX_ENHANCE_SEED)
|
||||
|
||||
def enhance(self, req: EnhancePromptRequest) -> EnhancePromptResponse:
|
||||
# Enhance never occupies the GPU slot (see PipelinesHandler.
|
||||
# evict_gpu_pipeline_for_prompt_enhancement) but still needs to mutually exclude with
|
||||
# generation and with itself — an abandoned/orphaned enhance call (e.g. the tab reloaded
|
||||
# mid-request) must not race a Generate click, a second Enhance click, or a generation
|
||||
# that's still loading its pipeline (reserved_generation_start covers that window; a bare
|
||||
# is_generation_running() check does not — see its own docstring). The "api" generation
|
||||
# slot gives us the mutual exclusion for free: it's the same bookkeeping every other
|
||||
# handler already does, and it doesn't require gpu_slot to be set.
|
||||
with self._generation.reserved_generation_start():
|
||||
gemma_root: str | None = None
|
||||
if req.provider == "local":
|
||||
gemma_root = self._text_handler.resolve_prompt_enhancer_root_if_downloaded()
|
||||
if gemma_root is None:
|
||||
raise HTTPError(409, "LOCAL_TEXT_ENCODER_NOT_AVAILABLE")
|
||||
elif not self.state.app_settings.gemini_api_key:
|
||||
raise HTTPError(400, "GEMINI_API_KEY_MISSING")
|
||||
|
||||
generation_id = uuid.uuid4().hex[:8]
|
||||
self._generation.start_api_generation(generation_id)
|
||||
try:
|
||||
enhanced = self._resolve_and_enhance(req, gemma_root)
|
||||
except HTTPError as e:
|
||||
self._generation.fail_generation(e.detail)
|
||||
raise
|
||||
except Exception as e:
|
||||
self._generation.fail_generation(str(e))
|
||||
raise HTTPError(500, str(e)) from e
|
||||
|
||||
self._generation.complete_generation(enhanced)
|
||||
return EnhancePromptResponse(enhancedPrompt=enhanced)
|
||||
|
||||
def _resolve_and_enhance(self, req: EnhancePromptRequest, gemma_root: str | None) -> str:
|
||||
if req.mediaType == "image":
|
||||
# No catalog LoRA concept for images (validated at the request level) — always an
|
||||
# explicit, image-domain system prompt, never the video-oriented generic fallback.
|
||||
system_prompt = (
|
||||
build_image_edit_system_prompt() if req.imagePath is not None
|
||||
else build_image_generation_system_prompt()
|
||||
)
|
||||
return self._run_free_rewrite(req, system_prompt, gemma_root)
|
||||
|
||||
if req.icLoraId is not None:
|
||||
ic_lora = self._lora_catalog_provider.get_ic_lora(req.icLoraId)
|
||||
if ic_lora is None:
|
||||
raise HTTPError(404, "LORA_CATALOG_ID_NOT_FOUND")
|
||||
return self._enhance_ic_lora(ic_lora, req, gemma_root)
|
||||
|
||||
if req.loraCatalogIds:
|
||||
loras: list[LoraCatalogItem] = []
|
||||
for catalog_id in req.loraCatalogIds:
|
||||
lora = self._lora_catalog_provider.get_lora(catalog_id)
|
||||
if lora is None:
|
||||
raise HTTPError(404, "LORA_CATALOG_ID_NOT_FOUND")
|
||||
loras.append(lora)
|
||||
return self._enhance_loras(loras, req, gemma_root)
|
||||
|
||||
if req.conditioningType is not None:
|
||||
system_prompt = build_conditioning_system_prompt(req.conditioningType)
|
||||
return self._run_free_rewrite(req, system_prompt, gemma_root)
|
||||
|
||||
return self._run_free_rewrite(req, self._default_video_system_prompt(req), gemma_root)
|
||||
|
||||
def _default_video_system_prompt(self, req: EnhancePromptRequest) -> str | None:
|
||||
if req.keyframes:
|
||||
return self._keyframe_system_prompt()
|
||||
return self._video_system_prompt(t2v=req.imagePath is None)
|
||||
|
||||
def _keyframe_system_prompt(self) -> str:
|
||||
spec = self._text_handler.active_ltx_model_spec()
|
||||
audio_visual = spec is not None and spec.wants_audio_visual_captions
|
||||
return build_keyframe_enhancement_system_prompt(audio_visual=audio_visual)
|
||||
|
||||
def _video_system_prompt(self, *, t2v: bool) -> str | None:
|
||||
"""The active model's own caption style, or None to keep each provider's default.
|
||||
|
||||
Only the audio-visual generations (2.5) need this: their captions cover the soundscape,
|
||||
which neither the generic Gemini fallback nor a 2.3-era prompt asks for.
|
||||
"""
|
||||
spec = self._text_handler.active_ltx_model_spec()
|
||||
if spec is None or not spec.wants_audio_visual_captions:
|
||||
return None
|
||||
return build_audio_visual_caption_system_prompt(t2v=t2v)
|
||||
|
||||
def enhance_for_generation(
|
||||
self,
|
||||
prompt: str,
|
||||
*,
|
||||
image_path: str | None,
|
||||
last_image_path: str | None = None,
|
||||
keyframes: list[KeyframeStill] | None = None,
|
||||
duration: int | None = None,
|
||||
fps: int | None = None,
|
||||
) -> str:
|
||||
"""Rewrite ``prompt`` on the local enhancer for a generation that's already started.
|
||||
|
||||
Only for the local text-encoding path: API encoding enhances server-side inside the
|
||||
same call, so it never gets here. Deliberately not `enhance()` — the caller already
|
||||
holds the generation slot, and there's no provider choice to make, only "is the
|
||||
enhancer on disk".
|
||||
|
||||
Never raises: enhancement is a quality step, so a missing checkpoint or a failed
|
||||
rewrite degrades to the prompt as typed rather than failing the generation.
|
||||
"""
|
||||
if not prompt.strip():
|
||||
return prompt
|
||||
gemma_root = self._text_handler.resolve_prompt_enhancer_root_if_downloaded()
|
||||
if gemma_root is None:
|
||||
logger.info("Skipping automatic enhancement: no local prompt enhancer downloaded")
|
||||
return prompt
|
||||
|
||||
try:
|
||||
pipeline = self._load_prompt_enhancer_pipeline(gemma_root)
|
||||
system_prompt = (
|
||||
self._keyframe_system_prompt()
|
||||
if keyframes
|
||||
else self._video_system_prompt(t2v=image_path is None)
|
||||
)
|
||||
seed = self._random_seed()
|
||||
if image_path is not None or keyframes:
|
||||
first_path = image_path or (keyframes[0][0] if keyframes else None)
|
||||
assert first_path is not None
|
||||
enhanced = pipeline.enhance_i2v(
|
||||
prompt,
|
||||
first_path,
|
||||
system_prompt=system_prompt,
|
||||
seed=seed,
|
||||
last_image_path=None if keyframes else last_image_path,
|
||||
keyframes=keyframes,
|
||||
duration=duration,
|
||||
fps=fps,
|
||||
)
|
||||
else:
|
||||
enhanced = pipeline.enhance_t2v(prompt, system_prompt=system_prompt, seed=seed)
|
||||
except Exception:
|
||||
logger.warning("Automatic local enhancement failed; using the prompt as typed", exc_info=True)
|
||||
return prompt
|
||||
|
||||
if not enhanced.strip():
|
||||
return prompt
|
||||
logger.info(
|
||||
"Enhanced prompt locally for generation (%d -> %d chars): %s",
|
||||
len(prompt),
|
||||
len(enhanced),
|
||||
enhanced,
|
||||
)
|
||||
return enhanced
|
||||
|
||||
def _enhance_loras(
|
||||
self, loras: list[LoraCatalogItem], req: EnhancePromptRequest, gemma_root: str | None
|
||||
) -> str:
|
||||
# None of the plain LoRAs in the catalog have a prompt_template today (only IC-LoRAs
|
||||
# do) — the multi-select path is always a free rewrite.
|
||||
system_prompt = build_lora_enhancement_system_prompt(loras)
|
||||
enhanced = self._run_free_rewrite(req, system_prompt, gemma_root)
|
||||
return enforce_trigger_placements(enhanced, loras)
|
||||
|
||||
def _enhance_ic_lora(
|
||||
self, ic_lora: IcLoraCatalogItem, req: EnhancePromptRequest, gemma_root: str | None
|
||||
) -> str:
|
||||
if ic_lora.prompt_template is not None:
|
||||
return self._run_template_fill(ic_lora, req, gemma_root)
|
||||
system_prompt = build_ic_lora_enhancement_system_prompt(ic_lora)
|
||||
enhanced = self._run_free_rewrite(req, system_prompt, gemma_root)
|
||||
return enforce_trigger_placements(enhanced, [ic_lora])
|
||||
|
||||
def _run_free_rewrite(
|
||||
self, req: EnhancePromptRequest, system_prompt: str | None, gemma_root: str | None
|
||||
) -> str:
|
||||
# Reject an invalid/unreadable/oversized path before it reaches either provider — the
|
||||
# API path in particular would otherwise base64-encode and ship arbitrary file bytes to
|
||||
# a third-party API with no gate at all.
|
||||
keyframes = self._validated_keyframes(req)
|
||||
image_path = (
|
||||
None if keyframes else normalize_optional_path(req.imagePath)
|
||||
)
|
||||
last_image_path = (
|
||||
None
|
||||
if req.mediaType == "image" or keyframes
|
||||
else normalize_optional_path(req.lastImagePath)
|
||||
)
|
||||
if image_path is not None:
|
||||
validate_image_file(image_path)
|
||||
if last_image_path is not None:
|
||||
validate_image_file(last_image_path)
|
||||
|
||||
seed = self._random_seed()
|
||||
first_path = image_path or (keyframes[0][0] if keyframes else None)
|
||||
if req.provider == "api":
|
||||
resolved_model = resolve_gemini_model(self.state.app_settings.gemini_model)
|
||||
logger.info("Enhancing prompt via Gemini API (%s)", resolved_model)
|
||||
api_key = self.state.app_settings.gemini_api_key
|
||||
if first_path is not None:
|
||||
return self._gemini_pipeline.enhance_i2v(
|
||||
req.prompt,
|
||||
first_path,
|
||||
system_prompt=system_prompt,
|
||||
seed=seed,
|
||||
api_key=api_key,
|
||||
model=resolved_model,
|
||||
last_image_path=last_image_path,
|
||||
keyframes=keyframes,
|
||||
duration=req.duration,
|
||||
fps=req.fps,
|
||||
)
|
||||
return self._gemini_pipeline.enhance_t2v(
|
||||
req.prompt,
|
||||
system_prompt=system_prompt,
|
||||
seed=seed,
|
||||
api_key=api_key,
|
||||
model=resolved_model,
|
||||
)
|
||||
|
||||
logger.info("Enhancing prompt via local Gemma")
|
||||
assert gemma_root is not None
|
||||
pipeline = self._load_prompt_enhancer_pipeline(gemma_root)
|
||||
if first_path is not None:
|
||||
return pipeline.enhance_i2v(
|
||||
req.prompt,
|
||||
first_path,
|
||||
system_prompt=system_prompt,
|
||||
seed=seed,
|
||||
last_image_path=last_image_path,
|
||||
keyframes=keyframes,
|
||||
duration=req.duration,
|
||||
fps=req.fps,
|
||||
)
|
||||
return pipeline.enhance_t2v(req.prompt, system_prompt=system_prompt, seed=seed)
|
||||
|
||||
def _validated_keyframes(self, req: EnhancePromptRequest) -> list[KeyframeStill] | None:
|
||||
if req.mediaType == "image" or not req.keyframes:
|
||||
return None
|
||||
frames: list[KeyframeStill] = []
|
||||
for keyframe in req.keyframes:
|
||||
path = normalize_optional_path(keyframe.imagePath)
|
||||
if path is None:
|
||||
raise HTTPError(400, "Each keyframe requires an image path")
|
||||
validate_image_file(path)
|
||||
frames.append((path, keyframe.frameIndex, keyframe.strength))
|
||||
frames.sort(key=lambda item: item[1])
|
||||
return frames
|
||||
|
||||
def _run_template_fill(
|
||||
self, ic_lora: IcLoraCatalogItem, req: EnhancePromptRequest, gemma_root: str | None
|
||||
) -> str:
|
||||
# req.imagePath is intentionally unused here — template fill is always a text-only
|
||||
# enhance_t2v call (the fixed template scaffold carries no reference-image slot), unlike
|
||||
# the free-rewrite IC-LoRA path below it, which does route an image through enhance_i2v.
|
||||
assert ic_lora.prompt_template is not None
|
||||
system_prompt = build_template_fill_system_prompt(ic_lora)
|
||||
seed = self._random_seed()
|
||||
try:
|
||||
if req.provider == "api":
|
||||
resolved_model = resolve_gemini_model(self.state.app_settings.gemini_model)
|
||||
logger.info("Enhancing prompt via Gemini API (%s)", resolved_model)
|
||||
raw = self._gemini_pipeline.enhance_t2v(
|
||||
req.prompt,
|
||||
system_prompt=system_prompt,
|
||||
seed=seed,
|
||||
api_key=self.state.app_settings.gemini_api_key,
|
||||
model=resolved_model,
|
||||
)
|
||||
else:
|
||||
logger.info("Enhancing prompt via local Gemma")
|
||||
assert gemma_root is not None
|
||||
pipeline = self._load_prompt_enhancer_pipeline(gemma_root)
|
||||
raw = pipeline.enhance_t2v(req.prompt, system_prompt=system_prompt, seed=seed)
|
||||
values = parse_template_fill_response(raw, set(ic_lora.prompt_template.placeholders))
|
||||
return fill_prompt_template(ic_lora.prompt_template, values)
|
||||
except ValueError as e:
|
||||
raise HTTPError(500, f"PROMPT_TEMPLATE_FILL_FAILED: {e}") from e
|
||||
|
||||
def _load_prompt_enhancer_pipeline(self, gemma_root: str) -> PromptEnhancerPipeline:
|
||||
self._pipelines.evict_gpu_pipeline_for_prompt_enhancement()
|
||||
device = os.getenv("LTX_PROMPT_ENHANCER_DEVICE", "").strip()
|
||||
if not device:
|
||||
device = get_device_type(self.config.device)
|
||||
logger.info("Loading local prompt enhancer on %s", device)
|
||||
return self._prompt_enhancer_pipeline_class.create(gemma_root, device)
|
||||
Reference in New Issue
Block a user