app.py

"""Mini Fooocus-inspired image generator using Flask and Stable Diffusion Turbo."""

from __future__ import annotations

import os
import threading
import uuid
from pathlib import Path
from typing import Dict, List

from flask import Flask, jsonify, render_template, request, url_for

try:
    import torch
    from diffusers import AutoPipelineForText2Image
except Exception as exc:  # pragma: no cover - surfaces missing deps quickly
    raise RuntimeError(
        "The required dependencies for the image generator are missing. "
        "Install them with 'pip install -r requirements.txt'."
    ) from exc


BASE_DIR = Path(__file__).resolve().parent
OUTPUT_DIR = BASE_DIR / "static" / "outputs"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


STYLE_PRESETS: Dict[str, Dict[str, str]] = {
    "none": {"label": "Original", "prompt": ""},
    "anime": {
        "label": "Anime",
        "prompt": "anime illustration, vibrant colors, crisp line art, dynamic lighting",
    },
    "cinematic": {
        "label": "Cinematic",
        "prompt": "cinematic composition, dramatic lighting, 35mm photograph, depth of field",
    },
    "realistic": {
        "label": "Realistic",
        "prompt": "photorealistic, ultra high definition, natural lighting, intricate detail",
    },
    "cyberpunk": {
        "label": "Cyberpunk",
        "prompt": "cyberpunk aesthetic, neon lights, futuristic cityscape, high contrast",
    },
    "three_d": {
        "label": "3D Render",
        "prompt": "ultra detailed 3d render, octane render, volumetric lighting, studio quality",
    },
}

MAX_IMAGES_PER_PROMPT = 8
RECENT_IMAGE_LIMIT = 20
MODEL_ID = "stabilityai/sd-turbo"


def _build_pipeline() -> AutoPipelineForText2Image:
    """Create and optimise the Stable Diffusion pipeline for fast inference."""

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    pipeline_kwargs = {"torch_dtype": dtype}
    if dtype == torch.float16:
        pipeline_kwargs["variant"] = "fp16"

    pipe = AutoPipelineForText2Image.from_pretrained(MODEL_ID, **pipeline_kwargs)

    pipe.to(device)
    pipe.set_progress_bar_config(disable=True)

    if device == "cuda":
        try:
            pipe.enable_xformers_memory_efficient_attention()
        except Exception:
            # xFormers is optional – continue without it if unavailable.
            pass

        try:
            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
        except Exception:
            # torch.compile may not be available or may fail on some setups.
            pass

    return pipe


PIPELINE_LOCK = threading.Lock()
PIPELINE = None


def get_pipeline():
    """Get or initialize the pipeline lazily."""
    global PIPELINE
    if PIPELINE is None:
        try:
            PIPELINE = _build_pipeline()
        except Exception as exc:
            raise RuntimeError(
                f"Failed to initialize pipeline: {exc}. "
                "Make sure dependencies are installed and you have internet access for model download."
            ) from exc
    return PIPELINE


app = Flask(__name__)


def list_recent_images(limit: int = RECENT_IMAGE_LIMIT) -> List[Dict[str, str]]:
    """Return the newest generated images for the gallery."""
    try:
        image_paths = sorted(
            OUTPUT_DIR.glob("*.png"),
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )[:limit]
        
        # Use Flask's url_for with app context
        with app.app_context():
            return [
                {
                    "url": url_for("static", filename=f"outputs/{path.name}"),
                    "filename": path.name,
                }
                for path in image_paths
            ]
    except Exception:
        # Return empty list if there are any issues
        return []


@app.route("/", methods=["GET"])
def index() -> str:
    presets = {key: data["label"] for key, data in STYLE_PRESETS.items()}
    recent_images = list_recent_images()
    return render_template(
        "index.html",
        presets=presets,
        recent_images=recent_images,
        image_options=[1, 2, 4, 6, 8],
    )


@app.route("/generate", methods=["POST"])
def generate_images():
    payload = request.get_json(force=True, silent=True) or {}
    prompt = (payload.get("prompt") or "").strip()
    style_key = payload.get("style", "none")
    count = int(payload.get("count", 1))

    if not prompt:
        return jsonify({"error": "Please enter a prompt."}), 400

    if style_key not in STYLE_PRESETS:
        style_key = "none"

    count = max(1, min(MAX_IMAGES_PER_PROMPT, count))

    style_prompt = STYLE_PRESETS[style_key]["prompt"]
    full_prompt = f"{prompt}, {style_prompt}" if style_prompt else prompt

    try:
        pipe = get_pipeline()
        with PIPELINE_LOCK:
            with torch.inference_mode():
                # SD Turbo doesn't support num_images_per_prompt well, generate one at a time
                generated_images = []
                for _ in range(count):
                    result = pipe(
                        prompt=full_prompt,
                        guidance_scale=0.0,
                        num_inference_steps=2,
                    )
                    # result.images is a list, even for single image
                    if isinstance(result.images, list):
                        generated_images.extend(result.images)
                    else:
                        generated_images.append(result.images)
    except Exception as exc:  # pragma: no cover - runtime failures only
        return (
            jsonify(
                {
                    "error": "Image generation failed. Please retry or check the console for details.",
                    "details": str(exc),
                }
            ),
            500,
        )

    saved_images: List[str] = []

    for image in generated_images:
        filename = f"{uuid.uuid4().hex}.png"
        image_path = OUTPUT_DIR / filename
        image.save(image_path)
        saved_images.append(url_for("static", filename=f"outputs/{filename}"))

    return jsonify({"images": saved_images})


if __name__ == "__main__":
    host = os.environ.get("FLASK_RUN_HOST", "0.0.0.0")
    port = int(os.environ.get("FLASK_RUN_PORT", 5000))
    debug = os.environ.get("FLASK_DEBUG", "0") == "1"
    app.run(host=host, port=port, debug=debug)