| import io |
| import os |
| import random |
| import threading |
| import time |
| from queue import Queue, Empty |
|
|
| import torch |
| from fastapi import FastAPI, Response |
| from fastapi.responses import JSONResponse, HTMLResponse |
| from fastapi.middleware.cors import CORSMiddleware |
| from diffusers import AutoPipelineForText2Image |
| from PIL import Image |
|
|
| MODEL_ID = os.getenv("MODEL_ID", "stabilityai/sd-turbo") |
|
|
| W = int(os.getenv("W", "384")) |
| H = int(os.getenv("H", "384")) |
|
|
| STEPS = int(os.getenv("STEPS", "3")) |
| GUIDANCE = float(os.getenv("GUIDANCE", "0.0")) |
|
|
| USE_FP32_ON_MPS = os.getenv("USE_FP32_ON_MPS", "1") == "1" |
| QUEUE_MAX = int(os.getenv("QUEUE_MAX", "2")) |
|
|
| |
| |
| |
| |
| P_MODERN = float(os.getenv("P_MODERN", "0.50")) |
| P_CLASSIC = float(os.getenv("P_CLASSIC", "0.30")) |
| P_EARLY = float(os.getenv("P_EARLY", "0.20")) |
|
|
| |
| CLASSIC_PEOPLE_WEIGHT = float(os.getenv("CLASSIC_PEOPLE_WEIGHT", "0.60")) |
| MODERN_PEOPLE_WEIGHT = float(os.getenv("MODERN_PEOPLE_WEIGHT", "0.35")) |
|
|
| |
| NEGATIVE = ( |
| |
| "frame, picture frame, painting frame, ornate frame, gold frame, " |
| "border, canvas edge, cropped canvas, mat, passepartout, " |
| "gallery wall, museum wall, hanging painting, framed artwork, " |
| "wood frame, gilded frame, edge of painting, " |
| |
| "text, watermark, logo, signature, letters, " |
| |
| "nsfw, nude, naked, porn, gore, violence, " |
| |
| "blur, blurry, out of focus, lowres, jpeg artifacts, " |
| |
| "bad anatomy, bad proportions, bad face, deformed face, " |
| "bad hands, malformed hands, deformed hands, " |
| "extra fingers, missing fingers, fused fingers, extra limbs, " |
| "hands in foreground, close-up hands, cropped hands, " |
| |
| "photorealistic, hyperrealistic, cgi, 3d render, plastic skin, anime, cartoon" |
| ) |
|
|
| |
|
|
| BASE_PAINTING_QUALITY = ( |
| "fine art painting, museum-quality artwork, painterly, expressive brushwork, " |
| "coherent artistic style, unified composition, natural color harmony, " |
| "sharp focus, high detail, canvas texture subtle, " |
| "no frame, no border, no canvas edge, " |
| "no photographic realism" |
| ) |
|
|
| COMPOSITION_CALM = ( |
| "balanced composition, calm pose, medium shot, hands not emphasized, " |
| "hands partially obscured by clothing or out of frame, " |
| "no dramatic hand gestures, hands not in foreground" |
| ) |
|
|
| LIGHTING_SOFT = "soft natural light, gentle contrast, pleasing tonal range" |
| LIGHTING_DRAMATIC = "dramatic chiaroscuro, deep shadows, warm highlights" |
|
|
| |
| EARLY_POOL = [ |
| "early medieval illuminated manuscript style, flat composition, symbolic forms, tempera, muted pigments", |
| "byzantine icon painting style, gold leaf tones, sacred atmosphere, stylized features, tempera on wood", |
| "gothic panel painting style, elongated forms, ornate patterns, flat background, tempera", |
| "romanese mural painting style, fresco texture, simplified figures, symbolic composition", |
| "medieval devotional painting style, stylized drapery, flat shapes, decorative borders implied (but no frame)", |
| ] |
|
|
| |
| CLASSIC_PEOPLE = [ |
| "early renaissance oil painting portrait, sfumato, subtle realism, classical balance", |
| "high renaissance portrait painting, refined anatomy, calm expression, old master", |
| "baroque oil painting figure scene, rich pigments, theatrical lighting", |
| "dutch golden age interior scene painting, soft window light, oil on canvas", |
| "romanticism portrait painting, warm skin tones, painterly texture, emotional mood", |
| ] |
|
|
| CLASSIC_LANDSCAPES = [ |
| "renaissance landscape painting, atmospheric perspective, classical composition", |
| "baroque landscape painting, dramatic sky, warm highlights, painterly", |
| "romantic landscape painting, luminous clouds, distant horizon, oil on canvas", |
| "classical pastoral landscape painting, soft light, calm mood, painterly", |
| ] |
|
|
| |
| MODERN_PEOPLE = [ |
| "impressionist portrait painting, visible brush strokes, light and color, soft edges", |
| "post-impressionist portrait painting, structured brushwork, rich color, painterly", |
| "fauvism portrait painting, bold color harmony, simplified shapes, expressive", |
| "expressionist figure painting, energetic brushwork, emotional color, painterly", |
| "modern figurative painting, simplified forms, contemporary palette, painterly", |
| ] |
|
|
| MODERN_LANDSCAPES = [ |
| "impressionist landscape painting, plein air, shimmering light, visible brush strokes", |
| "post-impressionist landscape painting, vibrant color, structured strokes, painterly", |
| "fauvism landscape painting, bold color fields, simplified forms, expressive", |
| "modern abstract landscape-inspired painting, color fields, texture, painterly", |
| "contemporary painting, abstract forms, subtle glitch-like texture, mixed media feel (still painterly)", |
| "minimal color-field painting, soft gradients, subtle texture, contemporary art", |
| ] |
|
|
| def weighted_choice(groups): |
| r = random.random() |
| acc = 0.0 |
| for p, name in groups: |
| acc += p |
| if r <= acc: |
| return name |
| return groups[-1][1] |
|
|
| def pick_epoch_group(): |
| total = P_MODERN + P_CLASSIC + P_EARLY |
| if total <= 0: |
| return "modern" |
| pm = P_MODERN / total |
| pc = P_CLASSIC / total |
| pe = P_EARLY / total |
| return weighted_choice([(pm, "modern"), (pc, "classic"), (pe, "early")]) |
|
|
| def pick_prompt(): |
| epoch = pick_epoch_group() |
|
|
| if epoch == "early": |
| style = random.choice(EARLY_POOL) |
| return f"{style}, fine art painting, no frame, no border, painterly" |
|
|
| if epoch == "classic": |
| style = random.choice(CLASSIC_PEOPLE) if random.random() < CLASSIC_PEOPLE_WEIGHT else random.choice(CLASSIC_LANDSCAPES) |
| return f"{style}, oil painting, museum quality, no frame, no border, calm pose" |
|
|
| style = random.choice(MODERN_PEOPLE) if random.random() < MODERN_PEOPLE_WEIGHT else random.choice(MODERN_LANDSCAPES) |
| return f"{style}, painterly, expressive brushwork, no frame, no border" |
|
|
| |
| |
| if random.random() < MODERN_PEOPLE_WEIGHT: |
| style = random.choice(MODERN_PEOPLE) |
| comp = COMPOSITION_CALM |
| else: |
| style = random.choice(MODERN_LANDSCAPES) |
| comp = "balanced composition, no frame, no border" |
| lighting = random.choice([LIGHTING_SOFT, "natural daylight, atmospheric light", "soft studio light"]) |
| return f"{style}, {BASE_PAINTING_QUALITY}, {lighting}, {comp}" |
|
|
| |
| if torch.backends.mps.is_available(): |
| DEVICE = "mps" |
| elif torch.cuda.is_available(): |
| DEVICE = "cuda" |
| else: |
| DEVICE = "cpu" |
|
|
| if DEVICE == "mps": |
| DTYPE = torch.float32 if USE_FP32_ON_MPS else torch.float16 |
| elif DEVICE == "cuda": |
| DTYPE = torch.float16 |
| else: |
| DTYPE = torch.float32 |
|
|
| app = FastAPI() |
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
| pipe = None |
| pipe_lock = threading.Lock() |
|
|
| q = Queue(maxsize=QUEUE_MAX) |
|
|
| latest_id = 0 |
| last_error = "" |
| last_gen_ms = None |
| generated_total = 0 |
| last_prompt = "" |
|
|
| def load_pipeline(): |
| global pipe |
| pipe = AutoPipelineForText2Image.from_pretrained( |
| MODEL_ID, |
| torch_dtype=DTYPE, |
| safety_checker=None, |
| feature_extractor=None, |
| ).to(DEVICE) |
| try: |
| pipe.set_progress_bar_config(disable=True) |
| except Exception: |
| pass |
|
|
| def render_png(): |
| global last_prompt |
| prompt = pick_prompt() |
| last_prompt = prompt |
|
|
| t0 = time.perf_counter() |
| with pipe_lock, torch.inference_mode(): |
| out = pipe( |
| prompt=prompt, |
| negative_prompt=NEGATIVE, |
| width=W, |
| height=H, |
| num_inference_steps=STEPS, |
| guidance_scale=GUIDANCE, |
| ) |
| img: Image.Image = out.images[0] |
| buf = io.BytesIO() |
| img.save(buf, format="PNG", optimize=True) |
| ms = (time.perf_counter() - t0) * 1000.0 |
| return buf.getvalue(), ms |
|
|
| def generator_loop(): |
| global latest_id, last_error, last_gen_ms, generated_total |
|
|
| while True: |
| try: |
| png, ms = render_png() |
|
|
| latest_id += 1 |
| last_gen_ms = ms |
| last_error = "" |
| generated_total += 1 |
|
|
| if q.full(): |
| try: |
| q.get_nowait() |
| except Empty: |
| pass |
|
|
| q.put((latest_id, png), timeout=1) |
|
|
| if DEVICE == "mps": |
| try: |
| torch.mps.empty_cache() |
| except Exception: |
| pass |
|
|
| except Exception as e: |
| last_error = repr(e) |
| time.sleep(0.5) |
|
|
| @app.on_event("startup") |
| async def startup(): |
| load_pipeline() |
| threading.Thread(target=generator_loop, daemon=True).start() |
|
|
| @app.get("/", response_class=HTMLResponse) |
| def root(): |
| try: |
| with open(os.path.join(os.path.dirname(__file__), "index.html"), "r", encoding="utf-8") as f: |
| return f.read() |
| except Exception: |
| return "<html><body style='background:black;color:white;font-family:system-ui'>index.html not found</body></html>" |
|
|
| @app.get("/health") |
| def health(): |
| return JSONResponse({ |
| "device": DEVICE, |
| "dtype": str(DTYPE), |
| "model": MODEL_ID, |
| "w": W, |
| "h": H, |
| "steps": STEPS, |
| "guidance": GUIDANCE, |
| "queue": q.qsize(), |
| "latest_id": latest_id, |
| "last_gen_ms": last_gen_ms, |
| "last_error": last_error, |
| "generated_total": generated_total, |
| "p_modern": P_MODERN, |
| "p_classic": P_CLASSIC, |
| "p_early": P_EARLY, |
| "classic_people_weight": CLASSIC_PEOPLE_WEIGHT, |
| "modern_people_weight": MODERN_PEOPLE_WEIGHT, |
| "last_prompt": last_prompt[:400], |
| }) |
|
|
| @app.get("/next") |
| def next_frame(): |
| fid, png = q.get(timeout=600) |
| return Response( |
| content=png, |
| media_type="image/png", |
| headers={"X-Frame-Id": str(fid), "Cache-Control": "no-store"}, |
| ) |
