Please enter the commit message for your changes. Lines starting
with '#' will be ignored, and an empty message aborts the commit. On branch main Initial commit Changes to be committed: new file: .DS_Store new file: .env new file: .gitignore new file: ai-worker/Dockerfile new file: ai-worker/requirements.txt new file: ai-worker/worker.py new file: background-worker/Dockerfile new file: background-worker/go.mod new file: background-worker/go.sum new file: background-worker/main.go new file: background-worker/market.go new file: background-worker/rmv.go new file: background-worker/rss.go new file: background-worker/sql_work.go new file: db/Dockerfile new file: db/init.sql new file: docker-compose.yml new file: server-app/dockerfile new file: server-app/go.mod new file: server-app/go.sum new file: server-app/main.go new file: volumes/.DS_Store new file: volumes/db-init/.DS_Store new file: volumes/db-init/data/news_rss_feeds.csv new file: volumes/web/.DS_Store new file: volumes/web/static/css/blog.css new file: volumes/web/static/css/index-lite.css new file: volumes/web/static/css/index.css new file: volumes/web/static/css/mandelbrot.css new file: volumes/web/static/img/minecraft.png new file: volumes/web/static/js/blog.js new file: volumes/web/static/js/index-lite.js new file: volumes/web/static/js/index.js new file: volumes/web/static/js/mandelbrot.js new file: volumes/web/static/media/cantina.mp3 new file: volumes/web/static/media/countdowns.json new file: volumes/web/static/media/gong.mp4 new file: volumes/web/template/blog.html new file: volumes/web/template/index-lite.html new file: volumes/web/template/index.html new file: volumes/web/template/mandelbrot.html
This commit is contained in:
hubble_dubble
183
ai-worker/worker.py
Normal file
183
ai-worker/worker.py
Normal file
@@ -0,0 +1,183 @@
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import psycopg2
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
|
||||
DB_URL = os.environ.get("DATABASE_URL", "postgres://app:appsecret@database:5432/appdb?sslmode=disable")
|
||||
DEFAULT_MODEL_PATH = Path(__file__).resolve().parent / "models" / "sd15"
|
||||
MODEL_PATH = Path(DEFAULT_MODEL_PATH)
|
||||
IMAGE_OUTPUT_DIR = Path(os.environ.get("IMAGE_OUTPUT_DIR") or "/app/images")
|
||||
IMAGE_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
MODEL_ID = "stabilityai/sd-turbo"
|
||||
WORKER_ID = os.environ.get("WORKER_ID") or f"{socket.gethostname()}:{os.getpid()}"
|
||||
BATCH_SIZE = int(os.environ.get("IMAGE_BATCH_SIZE") or "3")
|
||||
CLAIM_TTL_MINUTES = int(os.environ.get("IMAGE_CLAIM_TTL_MINUTES") or "20")
|
||||
pipe = None
|
||||
pipe_lock = threading.Lock()
|
||||
|
||||
|
||||
def connect():
|
||||
return psycopg2.connect(DB_URL)
|
||||
|
||||
|
||||
def load_pipeline() -> Optional[StableDiffusionPipeline]:
|
||||
global pipe
|
||||
|
||||
if pipe is not None:
|
||||
return pipe
|
||||
|
||||
with pipe_lock:
|
||||
if pipe is not None:
|
||||
return pipe
|
||||
|
||||
try:
|
||||
# Prüfen, ob lokal schon etwas im Modellordner liegt
|
||||
has_local_model = MODEL_PATH.is_dir() and any(MODEL_PATH.iterdir())
|
||||
|
||||
if has_local_model:
|
||||
logging.info("Lade Stable Diffusion Modell lokal aus %s", MODEL_PATH)
|
||||
pipe = StableDiffusionPipeline.from_pretrained(
|
||||
str(MODEL_PATH),
|
||||
torch_dtype=torch.float32,
|
||||
local_files_only=True,
|
||||
)
|
||||
else:
|
||||
logging.info("Kein lokales Modell gefunden – lade von Hugging Face (%s)", MODEL_ID)
|
||||
pipe = StableDiffusionPipeline.from_pretrained(
|
||||
MODEL_ID,
|
||||
torch_dtype=torch.float32
|
||||
)
|
||||
pipe.save_pretrained(MODEL_PATH)
|
||||
logging.info("Modell erfolgreich nach %s gespeichert", MODEL_PATH)
|
||||
|
||||
pipe = pipe.to("cpu")
|
||||
pipe.enable_attention_slicing()
|
||||
return pipe
|
||||
|
||||
except Exception as e:
|
||||
logging.error("Konnte Pipeline nicht laden: %s", e)
|
||||
pipe = None
|
||||
return None
|
||||
|
||||
|
||||
def fetch_articles_without_image(cur, limit=BATCH_SIZE):
|
||||
cur.execute(
|
||||
"""
|
||||
WITH claim AS (
|
||||
SELECT id
|
||||
FROM articles
|
||||
WHERE image IS NULL
|
||||
AND (
|
||||
image_claimed_at IS NULL
|
||||
OR image_claimed_at < now() - (%s * INTERVAL '1 minute')
|
||||
)
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE articles AS a
|
||||
SET image_claimed_at = now(),
|
||||
image_claimed_by = %s
|
||||
FROM claim
|
||||
WHERE a.id = claim.id
|
||||
RETURNING a.id, a.title, a.article_id
|
||||
""",
|
||||
(CLAIM_TTL_MINUTES, limit, WORKER_ID),
|
||||
)
|
||||
return cur.fetchall()
|
||||
|
||||
|
||||
def update_image(cur, article_id: int, data_uri: str):
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE articles
|
||||
SET image = %s,
|
||||
image_claimed_at = NULL,
|
||||
image_claimed_by = NULL
|
||||
WHERE id = %s
|
||||
""",
|
||||
(data_uri, article_id),
|
||||
)
|
||||
|
||||
|
||||
def release_claim(cur, article_id: int):
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE articles
|
||||
SET image_claimed_at = NULL,
|
||||
image_claimed_by = NULL
|
||||
WHERE id = %s AND image IS NULL
|
||||
""",
|
||||
(article_id,),
|
||||
)
|
||||
|
||||
|
||||
def safe_filename(name: str) -> str:
|
||||
safe = "".join(c if c.isalnum() or c in ("-", "_") else "_" for c in name)
|
||||
return safe or "image"
|
||||
|
||||
|
||||
def generate_image(prompt: str) -> Optional[tuple[str, bytes]]:
|
||||
model = load_pipeline()
|
||||
if model is None:
|
||||
return None
|
||||
try:
|
||||
img = model(
|
||||
prompt=prompt,
|
||||
num_inference_steps=8,
|
||||
guidance_scale=7,
|
||||
).images[0]
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
img_bytes = buf.getvalue()
|
||||
data = base64.b64encode(img_bytes).decode("ascii")
|
||||
return f"data:image/png;base64,{data}", img_bytes
|
||||
except Exception as e:
|
||||
logging.error("Bildgenerierung fehlgeschlagen: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
while True:
|
||||
try:
|
||||
with connect() as conn:
|
||||
with conn.cursor() as cur:
|
||||
rows = fetch_articles_without_image(cur)
|
||||
if not rows:
|
||||
logging.info("keine neuen Artikel ohne Bild")
|
||||
for aid, title, article_id in rows:
|
||||
prompt = title or "news illustration"
|
||||
result = generate_image(prompt)
|
||||
if result:
|
||||
data_uri, img_bytes = result
|
||||
filename = f"{safe_filename(article_id)}.png"
|
||||
out_path = IMAGE_OUTPUT_DIR / filename
|
||||
try:
|
||||
out_path.write_bytes(img_bytes)
|
||||
logging.info("Bild gespeichert unter %s", out_path)
|
||||
except Exception as e:
|
||||
logging.error("Konnte Bild nicht speichern (%s): %s", out_path, e)
|
||||
update_image(cur, aid, data_uri)
|
||||
logging.info("Bild gesetzt für Artikel %s", aid)
|
||||
else:
|
||||
release_claim(cur, aid)
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
logging.error("Fehler im Worker: %s", e)
|
||||
|
||||
time.sleep(30)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user