Spaces:
Running on Zero
Running on Zero
File size: 2,805 Bytes
b701455 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | import os
from typing import Optional
from src.Utilities import util
try:
import ollama # type: ignore
except ImportError: # pragma: no cover - optional dependency
ollama = None # type: ignore
DEFAULT_PREFIX = (
"masterpiece, best quality, (extremely detailed CG unity 8k wallpaper, "
"masterpiece, best quality, ultra-detailed, best shadow), high contrast, "
"(best illumination), ((cinematic light)), hyper detail, dramatic light, "
"depth of field"
)
def _resolve_prompt(source_prompt: Optional[str]) -> str:
try:
saved_prompt = util.load_parameters_from_file()[0]
except Exception:
saved_prompt = ""
return source_prompt if source_prompt else saved_prompt
def _extract_content(raw: str) -> str:
if "</think>" in raw:
raw = raw.split("</think>")[-1]
return raw.strip().strip('"')
def enhance_prompt(p: Optional[str]) -> str:
"""Enhance a prompt using an Ollama model when available."""
prompt = _resolve_prompt(p)
if not prompt or ollama is None:
if ollama is None:
print("Prompt enhancer skipped: Ollama client not available.")
return prompt
model = os.environ.get("PROMPT_ENHANCER_MODEL", "qwen3:0.6b").strip() or "qwen3:0.6b"
prefix = os.environ.get("PROMPT_ENHANCER_PREFIX", DEFAULT_PREFIX).strip()
try:
response = ollama.chat(
model=model,
messages=[
{
"role": "user",
"content": (
"You are a prompt composer for text-to-image diffusion models. "
"Rewrite the user prompt into a single comma-separated line "
"ordered by visual priority. Use short descriptive phrases "
"covering image type, mood, lighting, composition, background, "
"palette, key objects, artistic style, and subject appearance. "
"Wrap especially important clauses in parentheses. Avoid extra "
"commentary or labels. User prompt: "
f"{prompt}"
),
}
],
)
except Exception as exc: # pragma: no cover - network dependent
print(f"Prompt enhancer failed ({exc}); returning original prompt.")
return prompt
content = response.get("message", {}).get("content", "").strip()
if not content:
print("Prompt enhancer returned empty response; using original prompt.")
return prompt
enhanced = _extract_content(content)
if not enhanced:
return prompt
if prefix:
if not prefix.endswith(","):
prefix = prefix + ","
return f"{prefix} {enhanced}".strip()
return enhanced
|