Update app.py

app.py

@@ -4,7 +4,7 @@ import base64
 import requests
 from io import BytesIO
 from datetime import datetime
-from typing import Optional, Dict, Tuple
+from typing import Optional, Dict
 
 import gradio as gr
 from PIL import Image, ImageDraw, ImageFont
@@ -16,8 +16,7 @@ from huggingface_hub import InferenceClient
 HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
 hf = InferenceClient(token=HF_TOKEN)
 
-# IMPORTANT: only models that actually work via Inference API
-DEFAULT_MODELS = [
+MODELS = [
     "stabilityai/stable-diffusion-xl-base-1.0",
     "runwayml/stable-diffusion-v1-5",
 ]
@@ -44,61 +43,67 @@ def wrap_text(text, max_chars=60, max_lines=4):
 # ---------------------------
 # Core Generator
 # ---------------------------
-class GraphicsGenerator:
+class Generator:
     def __init__(self, client: InferenceClient):
         self.client = client
 
-    def analyze_repo(self, url: str) -> Tuple[Optional[dict], Optional[str]]:
-        try:
-            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)
-            if r.status_code != 200:
-                return None, r.text
-            j = r.json()
-            return {
-                "name": j["name"],
-                "description": j.get("description", ""),
-                "stars": j["stargazers_count"],
-                "forks": j["forks_count"],
-                "language": j.get("language", "Mixed"),
-            }, None
-        except Exception as e:
-            return None, str(e)
-
-    def build_prompt(self, data, style, colors):
-        return (
-            f"{style} GitHub repository banner. "
-            f"Project name: {data['name']}. "
-            f"{data['description']}. "
-            f"Color scheme: {colors}. "
-            f"Clean layout, no text artifacts, no watermark."
-        )
+    # ---- 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
 
-    def generate_ai_background(
+    # ---- AI ----
+    def generate_background(
         self,
-        prompt: str,
-        model: str,
-        width: int,
-        height: int,
+        prompt,
+        model,
+        width,
+        height,
+        negative=None,
         init_image: Optional[Image.Image] = None,
-    ) -> Image.Image:
+    ):
         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,
         )
 
-    def compose_final(self, bg: Image.Image, data: dict) -> Image.Image:
+    # ---- 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(
@@ -113,47 +118,139 @@ class GraphicsGenerator:
         draw.text((48, 40), data["name"], font=title, fill="white")
 
         y = 110
-        for line in wrap_text(data["description"]):
-            draw.text((48, y), line, font=body, fill="white")
-            y += 28
+        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")
 
-        footer = f"★ {data['stars']}   ⑂ {data['forks']}   {data['language']}"
-        draw.text((48, bg.height - 60), footer, font=body, fill="white")
         return bg
 
-generator = GraphicsGenerator(hf)
+
+gen = Generator(hf)
 
 # ---------------------------
-# Gradio
+# Handlers
 # ---------------------------
-def generate(repo_url, style, colors, model, width, height):
-    data, err = generator.analyze_repo(repo_url)
-    if err:
-        return None, err
+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
 
-    prompt = generator.build_prompt(data, style, colors)
-    bg = generator.generate_ai_background(prompt, model, width, height)
-    final = generator.compose_final(bg, data)
+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()
-    final.save(buf, format="PNG")
+    img.save(buf, format="PNG")
     link = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
-    return final, link
+    return img, link
 
-with gr.Blocks() as demo:
-    gr.Markdown("## GitHub Graphics Generator")
+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
 
-    repo = gr.Textbox(label="Repo URL")
-    style = gr.Dropdown(["Modern", "Minimal", "Professional"], value="Modern")
-    colors = gr.Textbox(value="Blue gradient")
-    model = gr.Dropdown(DEFAULT_MODELS, value=DEFAULT_MODELS[0])
-    width = gr.Slider(512, 1600, 1200, step=64)
-    height = gr.Slider(256, 1000, 628, step=64)
+# ---------------------------
+# 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")
 
-    btn = gr.Button("Generate")
-    img = gr.Image()
-    link = gr.Textbox(label="Download link")
+    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],
+            )
 
-    btn.click(generate, [repo, style, colors, model, width, height], [img, 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()