app.py

"""
Dispatch AI — UAE Heritage Art Generator
Themes: falconry, pearl diving, Bedouin life, modern Dubai, wind towers.
FLUX generates art, branded with Dispatch AI.
"""

import os
import io
import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image, ImageDraw, ImageFont

# --- Configuration -----------------------------------------------------------
HF_TOKEN = os.environ.get("HF_TOKEN", None)
MODEL_ID = "black-forest-labs/FLUX.1-schnell"

client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)

BG_COLOR = "#0A0F1A"
ACCENT = "#1FE0E6"

# Heritage themes
THEMES = {
    "Falconry": {
        "emoji": "🦅",
        "prompt": (
            "A majestic falcon perched on a gloved Emirati falconer's hand in the desert at golden hour, "
            "piercing eyes, intricate feather detail, warm sunset light, sand dunes, "
            "traditional Emirati heritage, ultra detailed wildlife photography, 4k"
        ),
        "description": "Falconry (Al Qanas) — the UAE's most revered traditional sport, practiced for centuries in the desert.",
    },
    "Pearl Diving": {
        "emoji": "🦪",
        "prompt": (
            "An ancient Emirati pearl diver in traditional dress underwater in the Arabian Gulf, "
            "holding an oyster shell with a glowing pearl, traditional dhow boat above, "
            "bioluminescent water, dreamlike ethereal lighting, orientalist painting, ultra detailed"
        ),
        "description": "Pearl diving (Al Ghaus) — the foundation of UAE's economy before oil, with divers holding their breath for minutes.",
    },
    "Bedouin Life": {
        "emoji": "⛺",
        "prompt": (
            "A traditional Bedouin camp in the Arabian desert at night, woven goat-hair tent, "
            "campfire glowing, camels resting, an elder telling stories, stars overhead, "
            "warm intimate atmosphere, orientalist painting style, ultra detailed"
        ),
        "description": "Bedouin life (Al Badu) — the nomadic desert lifestyle, with hospitality, poetry, and survival in the harsh environment.",
    },
    "Modern Dubai": {
        "emoji": "🏙️",
        "prompt": (
            "A breathtaking view of modern Dubai at sunset, Burj Khalifa towering above skyscrapers, "
            "Sheikh Zayed Road with luxury cars, golden and pink sky, palm trees, "
            "cinematic, ultra detailed cityscape photography, 4k"
        ),
        "description": "Modern Dubai — from fishing village to global metropolis in 50 years, the UAE's vision realized.",
    },
    "Wind Towers": {
        "emoji": "🏛️",
        "prompt": (
            "Traditional Arabian wind tower (barjeel) architecture in old Dubai's Bastakiya district, "
            "intricate coral-stone walls, narrow alleyways, warm golden afternoon sunlight, "
            "historic Emirati architecture, orientalist painting, ultra detailed"
        ),
        "description": "Wind towers (Barjeel) — ancient Emirati air conditioning, capturing breezes to cool homes before electricity.",
    },
    "Camel Caravan": {
        "emoji": "🐪",
        "prompt": (
            "A caravan of camels crossing the vast Empty Quarter desert at dawn, "
            "long shadows on rippled sand dunes, golden light, distant haze, "
            "epic scale, Lawrence of Arabia aesthetic, ultra detailed landscape, 4k"
        ),
        "description": "Camel caravans — the ships of the desert, carrying goods and people across the Arabian Peninsula for millennia.",
    },
    "Dhow Racing": {
        "emoji": "⛵",
        "prompt": (
            "Traditional Emirati dhow boats racing on Dubai Creek, white sails billowing, "
            "dramatic sky, splashing water, action shot, golden hour, "
            "ultra detailed maritime photography, 4k"
        ),
        "description": "Dhow racing — traditional wooden boats that carried the UAE's maritime trade, still raced in fierce competition today.",
    },
    "Date Palm Oasis": {
        "emoji": "🌴",
        "prompt": (
            "A lush date palm oasis in the UAE desert, green palm fronds against golden sand, "
            "traditional falaj water channel, sunlight filtering through palms, peaceful, "
            "ultra detailed landscape painting, orientalist style"
        ),
        "description": "Date palm oases (Al Nakheel) — the 'tree of life' sustaining desert communities, with sophisticated water systems.",
    },
    "Arabian Oryx": {
        "emoji": "🦌",
        "prompt": (
            "A majestic white Arabian oryx standing in the UAE desert at sunrise, "
            "elegant long horns, pristine white coat, golden sand dunes, "
            "wildlife photography, ultra detailed, 4k"
        ),
        "description": "Arabian Oryx — the UAE's national animal, saved from extinction through the late Sheikh Zayed's conservation program.",
    },
}

# Art styles
ART_STYLES = {
    "Photorealistic": "photorealistic, ultra detailed, 4k",
    "Oil Painting": "oil painting style, rich brushstrokes, orientalist art",
    "Watercolor": "delicate watercolor painting, soft washes, artistic",
    "Digital Art": "modern digital art, concept art style, vibrant colors",
    "Vintage Photograph": "vintage sepia photograph, aged, historical feel",
    "Cinematic": "cinematic, dramatic lighting, film still, movie poster style",
}

# Font for watermark
FONT_CANDIDATES = [
    "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
    "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
    "C:\\Windows\\Fonts\\arialbd.ttf",
    "arial.ttf",
]


def load_font(size):
    for path in FONT_CANDIDATES:
        if os.path.exists(path):
            return ImageFont.truetype(path, size)
    return ImageFont.load_default()


def add_watermark(img):
    """Add Dispatch AI watermark to bottom-right corner."""
    draw = ImageDraw.Draw(img)
    W, H = img.size
    font_size = max(16, int(W / 40))
    font = load_font(font_size)

    text = "Dispatch AI"
    # Get text bounding box
    bbox = draw.textbbox((0, 0), text, font=font)
    tw = bbox[2] - bbox[0]
    th = bbox[3] - bbox[1]

    x = W - tw - 20
    y = H - th - 20

    # Background bar
    draw.rectangle([x - 8, y - 4, x + tw + 8, y + th + 4], fill=BG_COLOR)

    # Outline
    for dx in (-1, 1):
        for dy in (-1, 1):
            draw.text((x + dx, y + dy), text, fill="black", font=font)
    # Text
    draw.text((x, y), text, fill=ACCENT, font=font)

    return img


def generate_heritage_art(theme, art_style, custom_prompt, add_brand):
    """Generate UAE heritage art with FLUX.1-schnell."""
    theme_data = THEMES.get(theme, THEMES["Falconry"])
    style_suffix = ART_STYLES.get(art_style, ART_STYLES["Photorealistic"])

    if custom_prompt and custom_prompt.strip():
        final_prompt = f"{custom_prompt.strip()}, {style_suffix}"
    else:
        final_prompt = f"{theme_data['prompt']}, {style_suffix}"

    w, h = 1024, 1024

    try:
        image = client.text_to_image(final_prompt, width=w, height=h)
        if not isinstance(image, Image.Image):
            image = Image.open(io.BytesIO(image)) if hasattr(image, "read") else Image.open(image)

        if add_brand:
            image = add_watermark(image)

        status = f"✅ Generated '{theme}' art in {art_style} style"
        if add_brand:
            status += " · Branded with Dispatch AI"

        info = f"### {theme_data['emoji']} {theme}\n\n{theme_data['description']}\n\n**Art Style:** {art_style}"

        return image, status, info
    except Exception as e:
        img = Image.new("RGB", (w, h), BG_COLOR)
        d = ImageDraw.Draw(img)
        d.text((w // 4, h // 2), f"Error: {str(e)[:50]}", fill=ACCENT)
        return img, f"❌ Error: {str(e)}", ""


# --- UI -----------------------------------------------------------------------
CSS = """
#dispatch-header h1 {
    color: #FFFFFF; font-size: 2.2rem; margin: 0;
    background: linear-gradient(90deg, #1FE0E6 0%, #FFFFFF 60%);
    -webkit-background-clip: text; -webkit-text-fill-color: transparent;
}
#dispatch-header p { color: #1FE0E6; font-size: 1.05rem; margin: 6px 0 0 0; }
.dispatch-footer { text-align: center; color: #8A8F9C; font-size: 0.9rem; padding-top: 8px; }
"""

with gr.Blocks(
    title="Dispatch AI — UAE Heritage Art Generator",
    theme=gr.themes.Base(
        primary_hue="cyan", secondary_hue="cyan", neutral_hue="slate",
        font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"],
    ).set(
        body_background_fill="#0A0F1A", body_background_fill_dark="#0A0F1A",
        body_text_color="#FFFFFF", body_text_color_dark="#FFFFFF",
        block_background_fill="#0E1424", block_background_fill_dark="#0E1424",
        block_border_color="#1FE0E6", block_border_width="1px",
        block_label_text_color="#1FE0E6", block_title_text_color="#1FE0E6",
        button_primary_background_fill="#1FE0E6", button_primary_background_fill_dark="#1FE0E6",
        button_primary_text_color="#0A0F1A", button_primary_border_color="#1FE0E6",
        input_background_fill="#0E1424", input_background_fill_dark="#0E1424",
        input_border_color="#1FE0E6", input_border_width="1px",
    ),
    css=CSS,
) as demo:
    with gr.Column(elem_id="dispatch-header"):
        gr.Markdown(
            """
            # Dispatch AI — UAE Heritage Art Generator
            Generate art inspired by UAE heritage · Falconry, Pearl Diving, Bedouin Life & more · Dispatch AI (FZE) · UAE
            """
        )

    with gr.Row():
        with gr.Column(scale=1):
            theme_select = gr.Dropdown(
                list(THEMES.keys()),
                label="Heritage Theme",
                value="Falconry",
                info="Choose a UAE heritage theme",
            )
            art_style_select = gr.Dropdown(
                list(ART_STYLES.keys()),
                label="Art Style",
                value="Photorealistic",
            )
            custom_prompt = gr.Textbox(
                label="Custom Prompt (optional — overrides theme prompt)",
                placeholder="Leave empty to use the theme's default prompt",
                lines=2,
            )
            add_brand_check = gr.Checkbox(
                label="Add 'Dispatch AI' watermark",
                value=True,
            )
            generate_btn = gr.Button("🎨 Generate Heritage Art", variant="primary")

        with gr.Column(scale=2):
            output_image = gr.Image(
                label="Generated Heritage Art",
                type="pil",
                show_download_button=True,
            )
            status_box = gr.Textbox(label="Status", interactive=False)
            theme_info = gr.Markdown()

    # Theme info updates
    def get_theme_info(theme):
        td = THEMES.get(theme, {})
        return f"### {td.get('emoji', '🎨')} {theme}\n\n{td.get('description', '')}"

    theme_select.change(fn=get_theme_info, inputs=theme_select, outputs=theme_info)
    # Initialize
    theme_info.value = get_theme_info("Falconry")

    # Events
    generate_btn.click(
        generate_heritage_art,
        inputs=[theme_select, art_style_select, custom_prompt, add_brand_check],
        outputs=[output_image, status_box, theme_info],
    )

    gr.Markdown(
        """
        <div class="dispatch-footer">
        © 2026 Dispatch AI (FZE) · Sharjah, UAE · License 10818 ·
        Model: FLUX.1-schnell · Celebrating UAE heritage through AI art
        </div>
        """
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch()