app.py

import os
import io
import base64
import numpy as np
import gradio as gr
import tempfile
from dotenv import load_dotenv
from huggingface_hub import InferenceClient
from PIL import Image, ImageDraw, ImageFilter, ImageFont


load_dotenv()
api_key = os.getenv("HF_API_TOKEN")

client = InferenceClient(api_key=api_key)

# Model
# IMAGE_MODEL = "black-forest-labs/FLUX.1-dev"
IMAGE_MODEL = "black-forest-labs/FLUX.1-schnell"  # faster, less quality

ODOO_COLORS = {
    "Odoo Purple  #714B67": {"hex": "#714B67", "name": "muted mauve purple"},
    "Odoo Teal    #1AD3BB": {"hex": "#1AD3BB", "name": "bright teal cyan"},
    "Sales Red    #E84646": {"hex": "#E84646", "name": "vivid coral red"},
    "Purchase Blue #3A97D4": {"hex": "#3A97D4", "name": "medium sky blue"},
    "HR Green     #1BBC9B": {"hex": "#1BBC9B", "name": "fresh emerald green"},
    "Finance Gold #F0A500": {"hex": "#F0A500", "name": "warm amber gold"},
    "Inventory    #00A09D": {"hex": "#00A09D", "name": "dark teal green"},
    "Project Blue #4C6EF5": {"hex": "#4C6EF5", "name": "bright indigo blue"},
    "Custom…":              {"hex": "custom",  "name": "custom"},
}

ODOO_MODULE_EXAMPLES = {
    "Sales":           "shopping cart, white glyph, bold minimal lines",
    "Purchase":        "shopping bag with arrow, clean outline",
    "Inventory":       "warehouse shelves, flat geometric",
    "Accounting":      "coins stack or ledger book",
    "HR":              "person silhouette with plus sign",
    "Project":         "kanban columns or gantt bars",
    "CRM":             "speech bubble with handshake",
    "Manufacturing":   "gear cog with wrench",
    "Website":         "globe with cursor",
    "Email Marketing": "envelope with lightning bolt",
    "Helpdesk":        "headset or life-ring",
    "Fleet":           "car front view, minimal",
    "Timesheets":      "clock with checkmark",
    "Expenses":        "receipt paper with dollar sign",
    "Sign":            "pen on document",
}

def make_transparent(img: Image.Image) -> Image.Image:
    img = img.convert("RGBA")
    data = np.array(img)
    white_mask = (
        (data[:, :, 0] > 245) &
        (data[:, :, 1] > 245) &
        (data[:, :, 2] > 245)
    )
    data[white_mask, 3] = 0
    return Image.fromarray(data)

def add_rounded_corners(img: Image.Image, radius: int = 60) -> Image.Image:
    img = img.convert("RGBA")
    mask = Image.new("L", img.size, 0)
    draw = ImageDraw.Draw(mask)
    draw.rounded_rectangle([0, 0, img.size[0], img.size[1]], radius=radius, fill=255)
    img.putalpha(mask)
    return img

def flatten_on_white(img: Image.Image) -> Image.Image:
    img = img.convert("RGBA")
    white_bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
    return Image.alpha_composite(white_bg, img).convert("RGB")

def hex_to_rgb(hex_color: str) -> tuple:
    h = hex_color.lstrip("#")
    return tuple(int(h[i:i+2], 16) for i in (0, 2, 4))

def add_rounded_corners(img: Image.Image, radius: int = 60) -> Image.Image:
    """Apply rounded corners with transparency (Odoo uses ~14px radius on 140px icons)."""
    img = img.convert("RGBA")
    mask = Image.new("L", img.size, 0)
    draw = ImageDraw.Draw(mask)
    draw.rounded_rectangle([0, 0, img.size[0], img.size[1]], radius=radius, fill=255)
    img.putalpha(mask)
    return img

def description_to_visual_elements(description: str) -> str:
    """
    Manually map module concepts to concrete visual primitives.
    This is the key step your current prompt skips entirely.
    """
    concept_map = {
        "invoice": "document sheet with folded corner",
        "purchase": "shopping cart or package box",
        "received": "downward arrow entering a box",
        "quantity": "stacked layers or count marks",
        "validation": "checkmark shield or lock",
        "3-way matching": "three overlapping circles or triangles meeting at center",
        "vendor bill": "stamped receipt",
        "stock": "warehouse shelf or stacked boxes",
        "billing": "paper with currency symbol",
        "prevent": "crossed-out shape or barrier",
    }
    found = []
    desc_lower = description.lower()
    for keyword, visual in concept_map.items():
        if keyword in desc_lower:
            found.append(visual)
    
    return "; ".join(found[:3])


def build_prompt(
    description: str,
    primary_color_name: str,
    style_hint: str,
    reference_module: str,
):
    visual_elements = description_to_visual_elements(description)
    if not visual_elements:
        visual_elements = description

    module_anchor = ""
    if reference_module and reference_module not in ("", "None"):
        module_anchor = f"Inspired by the visual language of the Odoo {reference_module} module icon."

    prompt = f"""
Odoo 18 enterprise application icon on pure white background.
Flat vector icon depicting: {visual_elements}.
Icon composition:
Two or three geometric shapes arranged to suggest the concept above.
Shapes overlap or nest to create visual hierarchy.
Negative space used deliberately as part of the design.
Rounded corners on all shapes (8–12px radius equivalent).
Color:
Primary fill: {primary_color_name}.
One or two complementary accent colors from the same tonal family.
No gradients. Flat fills only.
White used as internal cutout or separator.
Style:
Clean enterprise SaaS icon.
Same visual weight and abstraction level as Google Workspace or Notion icons.
Pixel-perfect at 128×128.
{module_anchor}
{style_hint}
Output format:
SVG-style vector graphic.
Transparent background (no square frame, no rounded app tile, no drop shadow).
No text, letters, numbers, or labels anywhere in the image.
No photorealism. No illustrations. No decorative borders.
""".strip()

    return " ".join(prompt.split())

def generate_icon(
    description: str,
    color_choice: str,
    custom_hex: str,
    style_hint: str,
    reference_module: str,
    reference_image,
    num_variants: int,
    apply_rounded: bool,
):
    if not description.strip():
        raise gr.Error("Please describe what the icon should represent.")

    color_data = ODOO_COLORS.get(color_choice, {"hex": "#714B67", "name": "muted mauve purple"})
    bg_hex  = color_data["hex"]
    bg_name = color_data["name"]

    if bg_hex == "custom":
        raw = custom_hex.strip()
        bg_hex = raw if raw.startswith("#") else f"#{raw}"
        bg_name = "custom color"

    prompt = build_prompt(
        description,
        bg_name,
        style_hint,
        reference_module,
    )

    if reference_image is not None:
        prompt += " Closely match the composition style of the uploaded reference icon."

    images = []
    png_files = []
    status_messages = [f"📝 Prompt: {prompt[:120]}…\n"]

    for i in range(int(1)):
        try:
            result = client.text_to_image(
                prompt=prompt,
                model=IMAGE_MODEL,
                width=512,
                height=512,
            )

            result = make_transparent(result)
            result = result.resize((140, 140), Image.LANCZOS)
            if apply_rounded:
                result = add_rounded_corners(result, radius=40)
            images.append(result)
            tmp = tempfile.NamedTemporaryFile(
                suffix=".png",
                delete=False
            )
            result.save(tmp.name, format="PNG")
            png_files.append(tmp.name)
            status_messages.append(
                f"✅ Variant {i+1} generated successfully"
            )
        except Exception as e:
            status_messages.append(f"❌ Variant {i+1} failed: {str(e)}")
            images.append(None)

    while len(images) < 3:
        images.append(None)
    while len(png_files) < 3:
        png_files.append(None)
    status = "\n".join(status_messages)
    return (
        images[0],
        png_files[0],
        prompt,
        status,
    )

CSS = """
.header-box {
    background: linear-gradient(135deg, #714B67 0%, #9B6B8F 100%);
    border-radius: 14px; padding: 22px 24px; margin-bottom: 18px;
}
.header-box h1 { color: white; margin: 0 0 4px; font-size: 1.65em; }
.header-box p  { color: #e8d5e3; margin: 0; font-size: 0.88em; }
.section-label {
    font-weight: 600; font-size: 0.85em; text-transform: uppercase;
    letter-spacing: 0.06em; color: #714B67; margin: 14px 0 4px;
}
.prompt-box textarea { font-family: monospace; font-size: 0.8em; background: #f8f5f7 !important; }
.fix-badge {
    background: #f0fdf4; border: 1px solid #86efac; border-radius: 8px;
    padding: 8px 12px; font-size: 0.82em; color: #166534;
}
"""

theme = gr.themes.Soft(primary_hue="purple")

with gr.Blocks(
    title="Odoo Icon Generator",
) as demo:
    gr.HTML("""
    <div class="header-box">
      <h1>🎨 Odoo Module Icon Generator</h1>
      <p>Generate flat material-design app icons using FLUX.1-dev · White glyph · Solid brand color · Odoo style</p>
    </div>
    """)

    with gr.Row(equal_height=False):
        with gr.Column(scale=1, min_width=320):
            gr.HTML('<div class="section-label">✏️ Icon concept</div>')
            description = gr.Textbox(
                label="What should the icon represent?",
                placeholder="e.g. a wrench crossed with a gear cog for a maintenance module",
                lines=2,
                show_label=True,
            )
            reference_module = gr.Dropdown(
                choices=["None"] + list(ODOO_MODULE_EXAMPLES.keys()),
                value="None",
                label="Reference Odoo Module"
            )
            gr.HTML('<div class="section-label">🎨 Background color</div>')
            color_choice = gr.Dropdown(
                choices=list(ODOO_COLORS.keys()),
                value="Odoo Purple  #714B67",
                label="Brand color preset",
            )
            custom_hex = gr.Textbox(
                label="Custom hex (active when 'Custom…' is selected)",
                placeholder="#2C7BE5",
                visible=False,
            )
            color_choice.change(
                fn=lambda c: gr.update(visible=(c == "Custom…")),
                inputs=color_choice,
                outputs=custom_hex,
            )
            gr.HTML('<div class="section-label">🛠️ Style options</div>')
            style_hint = gr.Textbox(
                label="Extra style hints (optional)",
                placeholder="e.g. rounded corners, thick stroke, ultra-minimalist",
                lines=1,
            )
            reference_image = gr.Image(
                label="📎 Upload a reference icon (optional)",
                type="pil",
                height=150,
            )
            num_variants = gr.Slider(
                minimum=1, 
                label="Number of variants",
                interactive=False
            )
            gr.HTML('<div class="section-label">🔧 Post-processing</div>')

            apply_rounded = gr.Checkbox(
                value=True, label="⬛ Rounded corners (Odoo style)",
                info="Adds ~14% corner radius matching Odoo icon shape",
            )
            generate_btn = gr.Button("🚀 Generate Icon", variant="primary", size="lg")

        with gr.Column(scale=1, min_width=320):

            gr.HTML('<div class="section-label">🖼️ Generated icons</div>')
            with gr.Row():
                out1 = gr.Image(label="Icon", height=200)

            with gr.Row():
                file1 = gr.File(label="Download PNG")

            gr.HTML('<div class="section-label">🔍 Prompt sent to model</div>')
            prompt_box = gr.Textbox(
                label="",
                lines=5,
                interactive=False,
                elem_classes=["prompt-box"],
                show_label=False,
            )

            status_box = gr.Textbox(
                label="Generation log",
                lines=4,
                interactive=False,
            )

    gr.HTML('<div class="section-label">💡 Quick examples</div>')
    gr.Examples(
        examples=[
            ["maintenance module — wrench crossed with a gear cog", "Odoo Purple  #714B67", "", "", "None", None, 3, True],
            ["delivery truck for logistics and shipping",            "Inventory    #00A09D", "", "bold thick outline", "None", None, 3, True],
        ],
        inputs=[description, color_choice, custom_hex, style_hint, reference_module, reference_image, num_variants, apply_rounded],
    )
    # ── Tips accordion ────────────────────────────────────────────────────
    with gr.Accordion("📖 Tips for great Odoo 18 icons", open=False):
        gr.Markdown("""
### Odoo 18 icon conventions
| Property | Value |
|----------|-------|
| Format | Square, 140×140 px (or 200×200) |
| Concept | Single glyph — one clear idea |
| Style | White icon on solid brand-color background |
| Effect | Flat + subtle drop shadow toward **bottom-left** |
| Strokes | Bold — icons must read at 48×48 px |
| Gradients | ❌ None |
| Text | ❌ None |

### Why does FLUX generate wrong colors?
FLUX.1 reads prompts as natural language — it doesn't parse hex codes like `#714B67`.
Instead it infers color from training data associations.
The **Enforce exact brand color** checkbox solves this by post-processing the image:
it detects non-white pixels and replaces them with your exact hex color after generation.

### Writing better prompts
- ✅ Be specific: `"invoice document with a euro sign"` beats `"accounting"`
- ✅ Name the style explicitly: `"flat, bold, white glyph, no gradients"`
- ✅ Reference a similar Odoo module to anchor the visual style
- ✅ Mention stroke weight: `"thick stroke"` vs `"thin elegant lines"`

### After generation
1. Download the PNG (right-click → Save image or use ⬇ button)
2. Resize to **140×140** or **200×200** px
3. Place at `your_module/static/description/icon.png`
4. Odoo Apps list will display it automatically ✅
        """)
    generate_btn.click(
        fn=generate_icon,
        inputs=[
            description,
            color_choice,
            custom_hex,
            style_hint,
            reference_image,
            num_variants,
            apply_rounded,
        ],
        outputs=[
            out1,
            file1,
            prompt_box,
            status_box,
        ]
    )
    # Footer
    gr.HTML("""
    <div style="
        margin-top: 30px;
        padding: 12px 16px;
        border-top: 1px solid #e5d7e3;
        text-align: center;
        font-size: 0.85em;
        color: #6b7280;
    ">
        🦥 Made with ❤️ by <b>Odoo Developer Mayna</b><br>
        <a href="https://github.com/stmayna"
           target="_blank"
           style="
               color: #714B67;
               text-decoration: none;
               font-weight: 600;
           ">
           🌿 GitHub
        </a>
    </div>
    """)

demo.launch(
    theme=theme,
    css=CSS,
    share=True,
)