app.py

# app.py
import os
import base64
import requests
from io import BytesIO
from datetime import datetime
from typing import Optional, Dict

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

# ---------------------------
# HF Client
# ---------------------------
HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
hf = InferenceClient(token=HF_TOKEN)

MODELS = [
    "stabilityai/stable-diffusion-xl-base-1.0",
    "runwayml/stable-diffusion-v1-5",
]

# ---------------------------
# Utilities
# ---------------------------
def wrap_text(text, max_chars=60, max_lines=4):
    words = text.split()
    lines, cur = [], []
    for w in words:
        test = " ".join(cur + [w])
        if len(test) <= max_chars:
            cur.append(w)
        else:
            lines.append(" ".join(cur))
            cur = [w]
        if len(lines) >= max_lines:
            break
    if cur and len(lines) < max_lines:
        lines.append(" ".join(cur))
    return lines[:max_lines]

# ---------------------------
# Core Generator
# ---------------------------
class Generator:
    def __init__(self, client: InferenceClient):
        self.client = client

    # ---- GitHub ----
    def fetch_repo(self, url: str):
        parts = url.replace("https://github.com/", "").split("/")
        owner, repo = parts[0], parts[1]
        r = requests.get(f"https://api.github.com/repos/{owner}/{repo}", timeout=10)
        j = r.json()
        return {
            "name": j["name"],
            "description": j.get("description") or "",
            "stars": j["stargazers_count"],
            "forks": j["forks_count"],
            "language": j.get("language") or "Mixed",
        }

    # ---- Prompt ----
    def build_prompt(self, base, style, colors, custom):
        p = f"{style} GitHub graphic. {base}. Color scheme: {colors}. Clean background."
        if custom:
            p += " " + custom
        return p

    # ---- AI ----
    def generate_background(
        self,
        prompt,
        model,
        width,
        height,
        negative=None,
        init_image: Optional[Image.Image] = None,
    ):
        if init_image:
            return self.client.image_to_image(
                image=init_image,
                prompt=prompt,
                negative_prompt=negative,
                model=model,
                width=width,
                height=height,
            )
        return self.client.text_to_image(
            prompt=prompt,
            negative_prompt=negative,
            model=model,
            width=width,
            height=height,
        )

    # ---- Layout ----
    def compose(
        self,
        bg: Image.Image,
        data: Dict,
        include_desc: bool,
        include_stats: bool,
        include_lang: bool,
    ):
        draw = ImageDraw.Draw(bg)
        try:
            title = ImageFont.truetype(
                "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 48
            )
            body = ImageFont.truetype(
                "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 20
            )
        except Exception:
            title = body = ImageFont.load_default()

        draw.text((48, 40), data["name"], font=title, fill="white")

        y = 110
        if include_desc and data.get("description"):
            for line in wrap_text(data["description"]):
                draw.text((48, y), line, font=body, fill="white")
                y += 28

        footer = []
        if include_stats:
            footer.append(f"★ {data['stars']}  ⑂ {data['forks']}")
        if include_lang:
            footer.append(data["language"])

        if footer:
            draw.text((48, bg.height - 60), "   ".join(footer), font=body, fill="white")

        return bg


gen = Generator(hf)

# ---------------------------
# Handlers
# ---------------------------
def handle_github(
    repo_url, style, colors, custom_prompt, negative,
    include_desc, include_stats, include_lang,
    model, width, height
):
    data = gen.fetch_repo(repo_url)
    base = f"{data['name']}. {data['description']}"
    prompt = gen.build_prompt(base, style, colors, custom_prompt)
    bg = gen.generate_background(prompt, model, width, height, negative)
    img = gen.compose(bg, data, include_desc, include_stats, include_lang)

    buf = BytesIO()
    img.save(buf, format="PNG")
    link = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
    return img, link

def handle_prompt(
    text, name, style, colors, custom_prompt, negative,
    include_desc, model, width, height
):
    data = {
        "name": name or "Custom Project",
        "description": text,
        "stars": 0,
        "forks": 0,
        "language": "N/A",
    }
    base = text
    prompt = gen.build_prompt(base, style, colors, custom_prompt)
    bg = gen.generate_background(prompt, model, width, height, negative)
    img = gen.compose(bg, data, include_desc, False, False)

    buf = BytesIO()
    img.save(buf, format="PNG")
    link = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
    return img, link

def handle_screenshot(
    image, style, colors, custom_prompt, negative,
    model, width, height
):
    prompt = gen.build_prompt("Application illustration", style, colors, custom_prompt)
    bg = gen.generate_background(prompt, model, width, height, negative, image)
    buf = BytesIO()
    bg.save(buf, format="PNG")
    link = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
    return bg, link

# ---------------------------
# UI
# ---------------------------
with gr.Blocks(title="GitHub Graphics Generator") as demo:
    gr.Markdown("## 🎨 GitHub Graphics Generator")

    model = gr.Dropdown(MODELS, value=MODELS[0], label="Model")
    width = gr.Slider(512, 1600, 1200, step=64, label="Width")
    height = gr.Slider(256, 1000, 628, step=64, label="Height")

    negative = gr.Textbox(label="Negative prompt", lines=2)
    custom_prompt = gr.Textbox(label="Custom prompt override", lines=2)

    with gr.Tabs():
        with gr.Tab("GitHub Repo"):
            repo = gr.Textbox(label="Repository URL")
            style = gr.Dropdown(["Modern", "Minimal", "Professional"], value="Modern")
            colors = gr.Textbox(value="Blue gradient")
            inc_desc = gr.Checkbox(True, label="Include description")
            inc_stats = gr.Checkbox(True, label="Include stats")
            inc_lang = gr.Checkbox(True, label="Include language")
            btn = gr.Button("Generate")
            out = gr.Image()
            link = gr.Textbox(label="Download link")

            btn.click(
                handle_github,
                [repo, style, colors, custom_prompt, negative,
                 inc_desc, inc_stats, inc_lang,
                 model, width, height],
                [out, link],
            )

        with gr.Tab("README / Prompt"):
            text = gr.Textbox(lines=8, label="Description or README")
            name = gr.Textbox(label="Project name")
            style2 = gr.Dropdown(["Modern", "Minimal", "Professional"], value="Modern")
            colors2 = gr.Textbox(value="Neutral")
            inc_desc2 = gr.Checkbox(True, label="Include text")
            btn2 = gr.Button("Generate")
            out2 = gr.Image()
            link2 = gr.Textbox(label="Download link")

            btn2.click(
                handle_prompt,
                [text, name, style2, colors2, custom_prompt, negative,
                 inc_desc2, model, width, height],
                [out2, link2],
            )

        with gr.Tab("Screenshot"):
            img = gr.Image(type="pil", label="Upload image")
            style3 = gr.Dropdown(["Modern", "Minimal", "Professional"], value="Modern")
            colors3 = gr.Textbox(value="Match image")
            btn3 = gr.Button("Generate")
            out3 = gr.Image()
            link3 = gr.Textbox(label="Download link")

            btn3.click(
                handle_screenshot,
                [img, style3, colors3, custom_prompt, negative,
                 model, width, height],
                [out3, link3],
            )

demo.launch()