app.py

import os
import random
import tempfile
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import gradio as gr
import spaces
import torch
from diffusers import StableDiffusionPipeline


MODEL_ID = os.getenv("MODEL_ID", "stabilityai/sd-turbo")
MAX_STEPS = int(os.getenv("MAX_STEPS", "12"))
GPU_DURATION = max(10, min(int(os.getenv("GPU_DURATION", "20")), 30))
BUILD_TAG = os.getenv("SPACE_BUILD_TAG", datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC"))

_pipe: Optional[StableDiffusionPipeline] = None
_pipe_cpu: Optional[StableDiffusionPipeline] = None

SAFE_RANDOM_PROMPTS = [
    "cinematic forest path at golden hour, ultra detailed, soft volumetric light",
    "futuristic neon city street in rain, reflections, moody atmosphere, photoreal",
    "cozy mountain cabin in winter, warm lights, snowfall, realistic texture",
    "epic fantasy castle on a cliff, dramatic sky, matte painting style",
    "macro photo of a dew-covered leaf, shallow depth of field, 8k",
    "anime style girl in cyberpunk tokyo, dynamic lighting, highly detailed",
    "astronaut walking on alien planet, cinematic composition, ultra wide shot",
    "vintage film noir detective in old city alley, monochrome, grain",
    "luxury sports car in studio, rim light, glossy reflections, product photo",
    "majestic dragon flying over snowy mountains, fantasy concept art",
]

ADULT_RANDOM_PROMPTS = [
    "tasteful sensual portrait, studio soft light, high detail skin, fine art",
    "boudoir fashion photo, cinematic shadows, elegant pose, photoreal",
    "romantic intimate couple portrait, warm light, shallow depth of field",
    "glamour portrait, beauty lighting, ultra realistic details, 85mm lens",
    "artistic nude silhouette, dramatic rim light, black background, fine art",
    "editorial lingerie photoshoot, luxury hotel room, cinematic color grading",
    "pin-up style portrait, retro lighting, detailed makeup, high contrast",
    "moody bedroom portrait, neon accent light, realistic skin texture",
    "high-fashion sensual photoshoot, magazine style, sharp focus",
    "tasteful body portrait, soft shadows, museum-grade fine art photography",
]

STYLE_TAGS = {
    "Cinematic": ["cinematic", "film", "volumetric", "dramatic"],
    "Photo": ["photo", "photoreal", "lens", "realistic", "macro", "studio soft light"],
    "Anime": ["anime", "manga"],
    "Fantasy": ["fantasy", "dragon", "castle", "matte painting", "alien"],
    "Studio": ["studio", "product", "rim light", "beauty lighting", "editorial", "fashion"],
}

STYLE_NEGATIVE_PROMPTS = {
    "Mix": "blurry, low quality, worst quality, jpeg artifacts, watermark",
    "Cinematic": "blurry, low quality, flat lighting, overexposed, underexposed, watermark",
    "Photo": "blurry, low quality, cartoon, cgi, plastic skin, overprocessed, watermark",
    "Anime": "blurry, low quality, bad anatomy, extra fingers, deformed face, watermark, text",
    "Fantasy": "blurry, low quality, low detail, muddy colors, watermark",
    "Studio": "blurry, low quality, bad skin texture, harsh shadows, watermark",
}

PHOTOREAL_PRO_NEGATIVE = (
    "blurry, low quality, worst quality, cartoon, anime, illustration, cgi,"
    " plastic skin, overprocessed skin, deformed face, extra fingers, watermark, text"
)

PORTRAIT_PRO_NEGATIVE = (
    "blurry, low quality, worst quality, cartoon, anime, illustration, cgi,"
    " deformed face, malformed face, asymmetrical eyes, extra eyes, duplicate face,"
    " multiple faces, twin face, cloned face, extra limbs, bad anatomy, watermark, text"
)

HUMAN_KEYWORDS = {
    "woman",
    "man",
    "girl",
    "boy",
    "female",
    "male",
    "person",
    "portrait",
    "face",
    "frau",
    "mann",
    "gesicht",
    "nackt",
    "nude",
}

ADULT_KEYWORDS = {
    "nude",
    "nudity",
    "nsfw",
    "explicit",
    "sex",
    "sexual",
    "erotic",
    "porn",
    "boobs",
    "breasts",
    "nipples",
    "penis",
    "vagina",
    "fetish",
    "lingerie",
    "naked",
}


def get_pipe(force_cpu: bool = False) -> StableDiffusionPipeline:
    global _pipe, _pipe_cpu

    use_gpu = torch.cuda.is_available() and not force_cpu
    if use_gpu:
        if _pipe is not None:
            return _pipe
        load_kwargs = {"torch_dtype": torch.float16, "variant": "fp16"}
        _pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, **load_kwargs)
        _pipe = _pipe.to("cuda")
        return _pipe

    if _pipe_cpu is not None:
        return _pipe_cpu
    load_kwargs = {"torch_dtype": torch.float32}
    _pipe_cpu = StableDiffusionPipeline.from_pretrained(MODEL_ID, **load_kwargs)
    _pipe_cpu = _pipe_cpu.to("cpu")
    return _pipe_cpu


def _generate_core(
    prompt: str,
    negative_prompt: str,
    steps: int,
    guidance: float,
    width: int,
    height: int,
    seed: str,
    adult_enabled: bool,
    realism_boost: bool,
    force_cpu: bool,
):
    if not prompt.strip():
        raise gr.Error("Prompt darf nicht leer sein.")

    prompt_text = prompt.strip()
    lowered_prompt = prompt_text.lower()
    if not adult_enabled and any(keyword in lowered_prompt for keyword in ADULT_KEYWORDS):
        raise gr.Error("Adult-Inhalte sind deaktiviert. Aktiviere den Adult-Schalter, um diesen Prompt zu nutzen.")

    pipe = get_pipe(force_cpu=force_cpu)
    steps = max(1, min(int(steps), MAX_STEPS))
    width = max(256, min(int(width), 1024))
    height = max(256, min(int(height), 1024))
    width = max(64, (width // 64) * 64)
    height = max(64, (height // 64) * 64)

    prompt_for_generation = prompt_text
    negative_for_generation = negative_prompt.strip() or ""
    lowered_prompt = prompt_for_generation.lower()
    is_human_prompt = any(keyword in lowered_prompt for keyword in HUMAN_KEYWORDS)

    if realism_boost:
        prompt_for_generation = (
            f"{prompt_for_generation}, photorealistic, ultra detailed, realistic lighting, natural skin texture,"
            " sharp focus, high dynamic range"
        )
        realism_negative = "cartoon, anime, illustration, lowres, deformed, oversaturated"
        negative_for_generation = f"{negative_for_generation}, {realism_negative}".strip(", ")

        if is_human_prompt:
            prompt_for_generation = (
                f"{prompt_for_generation}, single subject, one person, one face, centered composition,"
                " anatomically correct face, symmetrical eyes"
            )
            human_negative = (
                "multiple faces, duplicate face, extra eyes, extra nose, cloned face, bad facial anatomy"
            )
            negative_for_generation = f"{negative_for_generation}, {human_negative}".strip(", ")

    seed_text = "" if seed is None else str(seed).strip()
    seed_value = int(seed_text) if seed_text else 42
    generator_device = "cpu" if force_cpu else ("cuda" if torch.cuda.is_available() else "cpu")
    generator = torch.Generator(device=generator_device)
    generator.manual_seed(seed_value)

    result = pipe(
        prompt=prompt_for_generation,
        negative_prompt=negative_for_generation or None,
        num_inference_steps=steps,
        guidance_scale=float(guidance),
        width=width,
        height=height,
        generator=generator,
    )
    image = result.images[0]

    temp_dir = Path(tempfile.gettempdir()) / "pixelforge_downloads"
    temp_dir.mkdir(parents=True, exist_ok=True)
    download_path = temp_dir / f"pixelforge_{uuid.uuid4().hex[:12]}.png"
    image.save(download_path)

    return image, str(download_path)


@spaces.GPU(duration=GPU_DURATION)
def generate_image_gpu(
    prompt: str,
    negative_prompt: str,
    steps: int,
    guidance: float,
    width: int,
    height: int,
    seed: str,
    adult_enabled: bool,
    realism_boost: bool,
):
    return _generate_core(
        prompt,
        negative_prompt,
        steps,
        guidance,
        width,
        height,
        seed,
        adult_enabled,
        realism_boost,
        force_cpu=False,
    )


def generate_image(
    prompt: str,
    negative_prompt: str,
    steps: int,
    guidance: float,
    width: int,
    height: int,
    seed: str,
    adult_enabled: bool,
    realism_boost: bool,
):
    try:
        return generate_image_gpu(
            prompt,
            negative_prompt,
            steps,
            guidance,
            width,
            height,
            seed,
            adult_enabled,
            realism_boost,
        )
    except Exception as exc:
        message = str(exc)
        if "quota" in message.lower() or "zerogpu" in message.lower():
            return _generate_core(
                prompt,
                negative_prompt,
                steps,
                guidance,
                width,
                height,
                seed,
                adult_enabled,
                realism_boost,
                force_cpu=True,
            )
        if isinstance(exc, gr.Error):
            raise
        raise gr.Error(f"Generierung fehlgeschlagen: {exc}") from exc


def random_prompt(adult_enabled: bool, style: str):
    prompts = ADULT_RANDOM_PROMPTS if adult_enabled else SAFE_RANDOM_PROMPTS

    selected_style = style if style in STYLE_TAGS else "Mix"
    if selected_style == "Mix":
        filtered_prompts = prompts
    else:
        tags = STYLE_TAGS[selected_style]
        filtered_prompts = [
            prompt_text
            for prompt_text in prompts
            if any(tag in prompt_text.lower() for tag in tags)
        ]
        if not filtered_prompts:
            filtered_prompts = prompts

    selected_prompt = random.choice(filtered_prompts)

    negative_prompt = STYLE_NEGATIVE_PROMPTS.get(selected_style, STYLE_NEGATIVE_PROMPTS["Mix"])
    if adult_enabled:
        negative_prompt = f"{negative_prompt}, child, young, underage, teen, loli"

    return selected_prompt, negative_prompt


def apply_photoreal_pro_preset(adult_enabled: bool):
    base_negative = PHOTOREAL_PRO_NEGATIVE
    if adult_enabled:
        base_negative = f"{base_negative}, child, young, underage, teen, loli"

    return base_negative, 2.4, 896, 896, "Photo", True


def apply_portrait_pro_preset(adult_enabled: bool):
    base_negative = PORTRAIT_PRO_NEGATIVE
    if adult_enabled:
        base_negative = f"{base_negative}, child, young, underage, teen, loli"

    return base_negative, 2.8, 704, 960, "Photo", True


with gr.Blocks(title="PixelForge ZeroGPU") as demo:
    gr.Markdown("## PixelForge ZeroGPU")
    gr.Markdown("Leichte ZeroGPU-App für Text-zu-Bild mit SD-Turbo.")
    gr.Markdown(f"Build: {BUILD_TAG}")

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="Prompt", placeholder="z. B. cinematic cyberpunk city at night", lines=3)
            negative_prompt = gr.Textbox(label="Negative Prompt", value="blurry, low quality", lines=2)
            steps = gr.Slider(1, MAX_STEPS, value=min(8, MAX_STEPS), step=1, label="Steps")
            guidance = gr.Slider(0.0, 8.0, value=2.4, step=0.1, label="Guidance")
            width = gr.Slider(256, 1024, value=768, step=64, label="Bildbreite")
            height = gr.Slider(256, 1024, value=768, step=64, label="Bildhöhe")
            seed = gr.Textbox(label="Seed", value="42")
            adult_enabled = gr.Checkbox(label="Adult-Generierung erlauben (18+)", value=False)
            realism_boost = gr.Checkbox(label="Realismus Boost", value=True)
            style_select = gr.Dropdown(
                label="Random Style",
                choices=["Mix", "Cinematic", "Photo", "Anime", "Fantasy", "Studio"],
                value="Mix",
            )
            random_btn = gr.Button("Random Prompt", variant="secondary")
            photoreal_btn = gr.Button("Photoreal Pro Preset", variant="secondary")
            portrait_btn = gr.Button("Portrait Pro Preset", variant="secondary")
            run_btn = gr.Button("Bild erzeugen", variant="primary")

        with gr.Column(scale=1):
            image_out = gr.Image(label="Ergebnis", type="pil", elem_id="result-image", height=720)
            download_out = gr.File(label="Download", file_count="single")

    random_btn.click(
        fn=random_prompt,
        inputs=[adult_enabled, style_select],
        outputs=[prompt, negative_prompt],
        queue=False,
        api_name=False,
    )

    photoreal_btn.click(
        fn=apply_photoreal_pro_preset,
        inputs=[adult_enabled],
        outputs=[negative_prompt, guidance, width, height, style_select, realism_boost],
        queue=False,
        api_name=False,
    )

    portrait_btn.click(
        fn=apply_portrait_pro_preset,
        inputs=[adult_enabled],
        outputs=[negative_prompt, guidance, width, height, style_select, realism_boost],
        queue=False,
        api_name=False,
    )

    run_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance, width, height, seed, adult_enabled, realism_boost],
        outputs=[image_out, download_out],
        queue=False,
        api_name=False,
    )

    prompt.submit(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance, width, height, seed, adult_enabled, realism_boost],
        outputs=[image_out, download_out],
        queue=False,
        api_name=False,
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        show_error=True,
        show_api=False,
    )