Update app.py

app.py

@@ -1,51 +1,130 @@
 """
-Text-to-Game Generator - Groq Edition (Truly Free)
-- Code  : llama-3.1-8b-instant via Groq API (free, rate-limited per minute not monthly)
-- Images: canvas-drawn graphics only (zero API calls, zero quota, zero cost)
+Text-to-Game Generator - ZeroGPU Edition
+- Code  : Qwen/Qwen2.5-Coder-3B-Instruct   (transformers, ZeroGPU)
+- Images: deepseek-community/Janus-Pro-1B   (transformers, ZeroGPU)
 
-No HF token needed. No local models. No compilation. No OOM.
-Just add GROQ_API_KEY as a Space secret.
-
-Get a free Groq key (no credit card) at: console.groq.com
-Settings > Variables and secrets > New secret > Name: GROQ_API_KEY
+Space hardware: set to ZeroGPU in Space settings (free for public Spaces).
+Models load to CPU at startup, GPU is borrowed on-demand via @spaces.GPU.
+No API keys needed. No quota. Fully local.
 """
 
-import os
 import re
+import io
 import base64
 import traceback
 
+import torch
+import spaces
 import gradio as gr
-from openai import OpenAI
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import JanusForConditionalGeneration, JanusProcessor
+from PIL import Image
+
+# ---------------------------------------------------------------------------
+# Load code model at startup (CPU)
+# ---------------------------------------------------------------------------
+
+CODE_MODEL_ID = "Qwen/Qwen2.5-Coder-3B-Instruct"
+
+print("Loading Qwen2.5-Coder-3B tokenizer...")
+code_tokenizer = AutoTokenizer.from_pretrained(CODE_MODEL_ID)
+print("Loading Qwen2.5-Coder-3B model...")
+code_model = AutoModelForCausalLM.from_pretrained(
+    CODE_MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="cpu",
+)
+print("Code model ready.")
 
 # ---------------------------------------------------------------------------
-# Groq client
+# Load Janus-Pro-1B at startup (CPU)
 # ---------------------------------------------------------------------------
 
-CODE_MODEL   = "llama-3.1-8b-instant"
-GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
+IMAGE_MODEL_ID = "deepseek-community/Janus-Pro-1B"
 
+print("Loading Janus-Pro-1B processor...")
+janus_processor = JanusProcessor.from_pretrained(IMAGE_MODEL_ID)
+print("Loading Janus-Pro-1B model...")
+janus_model = JanusForConditionalGeneration.from_pretrained(
+    IMAGE_MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="cpu",
+)
+print("Image model ready.")
 
-def _check_key():
-    if not GROQ_API_KEY:
-        raise ValueError(
-            "GROQ_API_KEY is not set.\n"
-            "1. Sign up free (no credit card) at console.groq.com\n"
-            "2. Create an API key\n"
-            "3. Add it as a Space secret: Settings > Variables and secrets > Name: GROQ_API_KEY"
+# ---------------------------------------------------------------------------
+# Code generation — ZeroGPU
+# ---------------------------------------------------------------------------
+
+@spaces.GPU
+def run_code_model(prompt: str, max_new_tokens: int, temperature: float) -> str:
+    """Generate HTML game code using Qwen2.5-Coder-3B on GPU."""
+    device = "cuda"
+    code_model.to(device)
+
+    messages = [{"role": "user", "content": prompt}]
+    inputs = code_tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(device)
+
+    with torch.no_grad():
+        outputs = code_model.generate(
+            **inputs,
+            max_new_tokens=int(max_new_tokens),
+            temperature=float(temperature),
+            do_sample=True,
+            top_p=0.95,
         )
 
+    generated = outputs[0][inputs["input_ids"].shape[-1]:]
+    result = code_tokenizer.decode(generated, skip_special_tokens=True)
+    code_model.to("cpu")
+    torch.cuda.empty_cache()
+    return result
 
-def get_client():
-    _check_key()
-    return OpenAI(
-        base_url="https://api.groq.com/openai/v1",
-        api_key=GROQ_API_KEY,
-    )
+# ---------------------------------------------------------------------------
+# Image generation — ZeroGPU
+# ---------------------------------------------------------------------------
 
+@spaces.GPU
+def run_image_model(prompt: str) -> Image.Image:
+    """Generate a sprite image using Janus-Pro-1B on GPU."""
+    device = "cuda"
+    janus_model.to(device)
+
+    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
+    text_prompt = janus_processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = janus_processor(
+        text=text_prompt,
+        generation_mode="image",
+        return_tensors="pt",
+    ).to(device, dtype=torch.bfloat16)
+
+    janus_model.generation_config.num_return_sequences = 1
+
+    with torch.no_grad():
+        outputs = janus_model.generate(
+            **inputs,
+            generation_mode="image",
+            do_sample=True,
+            use_cache=True,
+        )
+
+    decoded = janus_model.decode_image_tokens(outputs)
+    images = janus_processor.postprocess(
+        list(decoded.float()), return_tensors="PIL.Image.Image"
+    )
+    pil_img = images["pixel_values"][0]
+    janus_model.to("cpu")
+    torch.cuda.empty_cache()
+    return pil_img
 
 # ---------------------------------------------------------------------------
-# Game type configs — all use canvas drawing, no image files
+# Game type configs
 # ---------------------------------------------------------------------------
 
 GAME_TYPES = {
@@ -60,9 +139,10 @@ GAME_TYPES = {
             "- At least 5 platforms, 2 moving enemies, and a goal to reach.\n"
             "- Show score/lives on the canvas.\n"
             "- Use requestAnimationFrame for the game loop.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors (fillRect, arc, fillStyle, etc).\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
+            "- For every image asset use: const img = new Image(); img.src = 'sprite_NAME.png';\n"
+            "  e.g. sprite_player.png, sprite_enemy.png, sprite_background.png. Max 3 sprites.\n"
+            "- No external libraries or CDN links.\n"
+            "- Theme all visuals to match: {theme}.\n"
             "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
@@ -78,9 +158,9 @@ GAME_TYPES = {
             "- Display health, score, and wave on canvas.\n"
             "- Game-over screen with score and restart button.\n"
             "- Use requestAnimationFrame for the game loop.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors only.\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
+            "- For every image asset: const img = new Image(); img.src = 'sprite_NAME.png'; Max 3 sprites.\n"
+            "- No external libraries or CDN links.\n"
+            "- Theme everything to match: {theme}.\n"
             "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
@@ -95,9 +175,9 @@ GAME_TYPES = {
             "- At least 15x10 tile grid, collectible keys, and a locked exit.\n"
             "- Show a move counter or timer.\n"
             "- Win screen when player collects all keys and exits.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors only.\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
+            "- For every image asset: const img = new Image(); img.src = 'sprite_NAME.png'; Max 3 sprites.\n"
+            "- No external libraries or CDN links.\n"
+            "- Theme everything to match: {theme}.\n"
             "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
@@ -112,9 +192,9 @@ GAME_TYPES = {
             "- Collision ends the game; show time survived as score.\n"
             "- High score stored in a JS variable; restart button on game-over screen.\n"
             "- Use requestAnimationFrame for the game loop.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors only.\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
+            "- For every image asset: const img = new Image(); img.src = 'sprite_NAME.png'; Max 3 sprites.\n"
+            "- No external libraries or CDN links.\n"
+            "- Theme everything to match: {theme}.\n"
             "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
@@ -130,9 +210,9 @@ GAME_TYPES = {
             "- Clear win/lose conditions and a score display.\n"
             "- Game-over / win screen with a restart button.\n"
             "- Use requestAnimationFrame for the game loop.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors only.\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
+            "- For every image asset: const img = new Image(); img.src = 'sprite_NAME.png'; Max 3 sprites.\n"
+            "- No external libraries or CDN links.\n"
+            "- Theme everything vividly to match: {theme}.\n"
             "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
@@ -149,7 +229,66 @@ THEME_EXAMPLES = {
 }
 
 # ---------------------------------------------------------------------------
-# Code generation
+# Sprite helpers
+# ---------------------------------------------------------------------------
+
+def _pil_to_data_uri(pil_image: Image.Image, size: int = 64) -> str:
+    img = pil_image.resize((size, size), Image.LANCZOS)
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+
+
+def _colored_placeholder(name: str) -> str:
+    colours = ["#e74c3c", "#3498db", "#2ecc71", "#f39c12", "#9b59b6", "#1abc9c"]
+    colour  = colours[abs(hash(name)) % len(colours)]
+    label   = name.replace("sprite_", "")[:6]
+    svg = (
+        '<svg xmlns="http://www.w3.org/2000/svg" width="64" height="64">'
+        '<rect width="64" height="64" fill="' + colour + '" rx="8"/>'
+        '<text x="32" y="38" font-size="10" text-anchor="middle" fill="white">' + label + '</text>'
+        '</svg>'
+    )
+    return "data:image/svg+xml;base64," + base64.b64encode(svg.encode()).decode()
+
+
+def _sprite_prompt(sprite_name: str, theme: str) -> str:
+    label = sprite_name.replace("sprite_", "").replace("_", " ")
+    return (
+        "Pixel-art game sprite of a " + label + " for a " + theme + " themed game. "
+        "Vibrant colours, clear silhouette, plain dark background, 64x64 style."
+    )
+
+
+def _find_sprite_filenames(html_code: str) -> list:
+    pattern = r"""(?:src\s*=\s*['"]|\.src\s*=\s*['"])\s*(sprite_[a-zA-Z0-9_]+\.png)\s*['"]"""
+    return list(dict.fromkeys(re.findall(pattern, html_code)))[:3]
+
+
+def _inject_sprites(html_code: str, sprite_map: dict) -> str:
+    for fname, data_uri in sprite_map.items():
+        html_code = html_code.replace('"' + fname + '"', '"' + data_uri + '"')
+        html_code = html_code.replace("'" + fname + "'", "'" + data_uri + "'")
+    return html_code
+
+
+def generate_sprites(sprite_names: list, theme: str) -> dict:
+    if not sprite_names:
+        return {}
+    mapping = {}
+    for fname in sprite_names:
+        name_no_ext = fname.replace(".png", "")
+        prompt = _sprite_prompt(name_no_ext, theme)
+        try:
+            pil_img = run_image_model(prompt)
+            mapping[fname] = _pil_to_data_uri(pil_img, size=64)
+        except Exception as exc:
+            print("[Janus] Failed '" + fname + "': " + str(exc))
+            mapping[fname] = _colored_placeholder(name_no_ext)
+    return mapping
+
+# ---------------------------------------------------------------------------
+# Full generation pipeline
 # ---------------------------------------------------------------------------
 
 def generate_game_code(game_type: str, theme: str, temperature: float, max_new_tokens: int):
@@ -160,43 +299,41 @@ def generate_game_code(game_type: str, theme: str, temperature: float, max_new_t
     user_prompt = template.format(theme=theme.strip())
     system_msg  = (
         "You are an expert HTML5 game developer. "
-        "Write a complete, working, single-file HTML5 game using only canvas 2D drawing. "
-        "Never use image files or external resources. "
-        "Draw everything with canvas shapes, colors, and text. "
+        "Write a complete, working, single-file HTML5 game using canvas and vanilla JavaScript. "
+        "Use new Image() with src='sprite_NAME.png' for every visual asset (max 3 sprites). "
         "Output ONLY the raw HTML - no markdown fences, no explanation."
     )
+    full_prompt = system_msg + "\n\n" + user_prompt
 
     try:
-        client = get_client()
-        completion = client.chat.completions.create(
-            model=CODE_MODEL,
-            messages=[
-                {"role": "system", "content": system_msg},
-                {"role": "user",   "content": user_prompt},
-            ],
-            max_tokens=int(max_new_tokens),
-            temperature=float(temperature),
-        )
-        code = completion.choices[0].message.content.strip()
+        # Step 1: Generate HTML via Qwen2.5-Coder-3B
+        code = run_code_model(full_prompt, max_new_tokens, temperature)
         code = re.sub(r"^```[a-zA-Z]*\n?", "", code).strip()
         code = re.sub(r"\n?```$",           "", code).strip()
         if "<html" not in code.lower() and "<!doctype" not in code.lower():
             code = _wrap_in_html(code, theme)
 
-        status = "Done! Game generated. Click Launch Game to play."
-        return code, status, _build_preview(code)
+        # Step 2: Detect sprites
+        sprite_names = _find_sprite_filenames(code)
+
+        # Step 3: Generate sprites via Janus-Pro-1B
+        sprite_map = generate_sprites(sprite_names, theme.strip())
+
+        # Step 4: Inject base64 images
+        final_code = _inject_sprites(code, sprite_map)
+
+        n_real     = sum(1 for v in sprite_map.values() if "image/png" in v)
+        n_fallback = len(sprite_map) - n_real
+        status = (
+            "Done! " + str(len(sprite_names)) + " sprite(s) - "
+            + str(n_real) + " by Janus-Pro-1B, "
+            + str(n_fallback) + " fallback. Click Launch Game to play."
+        )
+        return final_code, status, _build_preview(final_code)
 
     except Exception as exc:
         traceback.print_exc()
-        err = str(exc)
-        if "invalid_api_key" in err.lower() or "401" in err:
-            err = "Invalid GROQ_API_KEY. Check your key at console.groq.com."
-        elif "429" in err or "rate" in err.lower():
-            err = "Rate limited by Groq - wait a few seconds and try again."
-        else:
-            err = "Error: " + str(exc)
-        return "", err, _placeholder_html(err)
-
+        return "", "Error: " + str(exc), _placeholder_html("Error: " + str(exc))
 
 # ---------------------------------------------------------------------------
 # HTML helpers
@@ -237,7 +374,6 @@ def launch_game(code: str) -> str:
         return _placeholder_html("No game code yet - generate a game first.")
     return _build_preview(code)
 
-
 # ---------------------------------------------------------------------------
 # UI helpers
 # ---------------------------------------------------------------------------
@@ -249,21 +385,21 @@ def update_type_description(game_type: str) -> str:
 def get_first_theme(game_type: str) -> str:
     return THEME_EXAMPLES[game_type][0][0]
 
-
 # ---------------------------------------------------------------------------
 # Gradio UI
 # ---------------------------------------------------------------------------
 
 def build_ui():
-    with gr.Blocks(title="Game Generator - Free") as demo:
+    with gr.Blocks(title="Game Generator - ZeroGPU") as demo:
 
         gr.Markdown(
-            "# Game Generator (Truly Free)\n"
-            "**Llama-3.1-8B** via Groq generates complete playable games instantly.\n"
-            "All graphics are drawn with canvas — no image API, no quota, no cost.\n\n"
-            "> **Setup:** Sign up free at [console.groq.com](https://console.groq.com) "
-            "(no credit card needed), create an API key, then add it as a Space secret "
-            "named `GROQ_API_KEY`."
+            "# Game Generator (ZeroGPU)\n"
+            "**Qwen2.5-Coder-3B** generates the game code. "
+            "**Janus-Pro-1B** generates every sprite. "
+            "Both run locally on ZeroGPU — no API keys, no quota.\n\n"
+            "> Make sure your Space hardware is set to **ZeroGPU** "
+            "(Space Settings → Hardware → ZeroGPU). "
+            "First run downloads models (~9GB total) — subsequent runs are faster."
         )
 
         with gr.Row():
@@ -297,11 +433,11 @@ def build_ui():
                     label="Temperature - higher = more creative",
                 )
                 max_tokens_slider = gr.Slider(
-                    minimum=1000, maximum=6000, value=4000, step=500,
-                    label="Max tokens - more = longer / more complete game",
+                    minimum=1000, maximum=4000, value=2500, step=500,
+                    label="Max tokens - more = longer game",
                 )
 
-                generate_btn = gr.Button("Generate Game", variant="primary")
+                generate_btn = gr.Button("Generate Game + Sprites", variant="primary")
                 gen_status   = gr.Markdown(value="_No game generated yet._")
 
             with gr.Column(scale=2, min_width=500):
@@ -309,7 +445,7 @@ def build_ui():
                 gr.Markdown("## 3. Generated code (editable)")
 
                 code_box = gr.Code(
-                    label="HTML source",
+                    label="HTML source (sprites embedded as base64)",
                     language="html",
                     lines=12,
                     interactive=True,
@@ -336,15 +472,14 @@ def build_ui():
 
         gr.Markdown(
             "---\n"
-            "**Why Groq is truly free:** Groq uses per-minute rate limits, not monthly credit quotas. "
-            "If you hit a limit, wait 60 seconds — not 30 days. "
-            "Games use only canvas drawing so no image API is ever called. "
+            "**Code:** Qwen2.5-Coder-3B-Instruct (~6.2GB) via transformers + ZeroGPU. "
+            "**Images:** Janus-Pro-1B (~3.8GB) via transformers + ZeroGPU. "
+            "GPU is borrowed only during generation then released — no idle VRAM usage. "
             "Edit the HTML and click **Launch Game** to hot-reload."
         )
 
     return demo
 
-
 # ---------------------------------------------------------------------------
 # Entry point
 # ---------------------------------------------------------------------------