diffusiongemma-codegen-8192

Running on Zero

App Files Files Community

mfielding92 commited on Jun 23

Commit

e768b88

verified ·

1 Parent(s): 57eee12

Update app.py

Browse files

Files changed (1) hide show

app.py +52 -30

app.py CHANGED Viewed

@@ -1,3 +1,4 @@
 """
 Gemma Diffusion — live website builder (gradio.Server backend + custom frontend).
@@ -77,7 +78,7 @@ MAX_ITERS_CAP = 120  # hard cap on denoising steps per block
 # ZeroGPU: the 26B checkpoint (~49 GB bf16) needs the full backing card.
 GPU_SIZE = os.environ.get("GDIFF_GPU_SIZE", "xlarge")
-SYSTEM_PROMPT = (
     "You are an expert front-end web developer with great visual taste. When asked to "
     "build or change a web page, respond with a SINGLE, complete, self-contained HTML5 "
     "document. Put all CSS in a <style> tag and any JavaScript in a <script> tag inside "
@@ -87,10 +88,22 @@ SYSTEM_PROMPT = (
     "<!DOCTYPE html>."
 )
 _MARKER_RE = re.compile(
     r"<\|?(?:channel|turn|think|image|audio|video|tool(?:_call|_response)?)\|?>"
 )
-_FENCE_RE = re.compile(r"```(?:html)?\s*(.*?)\s*```", re.DOTALL)
 # --------------------------------------------------------------------------- #
@@ -110,24 +123,26 @@ CANVAS_LEN = model.config.canvas_length
 PAD_ID = tokenizer.pad_token_id or 0
 print(f"[gdiff] Model ready | canvas_length={CANVAS_LEN}", flush=True)
-# Cache of the last *cleaned* page so a follow-up tweak can warm-start in place.
 model._last_clean_html = None
 # --------------------------------------------------------------------------- #
 # Helpers (CPU-only; safe to run in the gradio.Server main process)
 # --------------------------------------------------------------------------- #
-def warm_canvas_from_cache():
-    """Starting canvas (first block) built from the previous *cleaned* page.
     Returns a CPU tensor (it is pickled across the ZeroGPU process boundary and moved
-    to CUDA inside the GPU worker). We re-tokenize the cleaned HTML rather than reuse
-    raw output tokens so a mangled header can't compound across tweaks.
     """
-    html = getattr(model, "_last_clean_html", None)
-    if not html:
         return None
-    ids = tokenizer(html, add_special_tokens=False).input_ids[:CANVAS_LEN]
     if not ids:
         return None
     if len(ids) < CANVAS_LEN:
@@ -151,12 +166,7 @@ def clean_text(text: str) -> str:
 def extract_html(text: str) -> str:
-    """Pull a usable HTML document out of the (possibly mangled) model output.
-    Anchor on the first intact structural tag and rebuild whatever the diffused tweak ate
-    off the front, so the result is always a valid document (never quirks mode and never a
-    broken ``DOCTYPE>`` / ``html lang=`` header).
-    """
     text = clean_text(text)
     fenced = _FENCE_RE.search(text)
     if fenced:
@@ -181,6 +191,15 @@ def extract_html(text: str) -> str:
     return text.strip()
 class QueueDiffusionStreamer(BaseStreamer):
     def __init__(self, tok, q: "queue_lib.Queue"):
         self.tok = tok
@@ -212,12 +231,12 @@ class QueueDiffusionStreamer(BaseStreamer):
         self.q.put(("end", self._decode(self.confirmed_ids), self.block, self.step))
-def build_messages(history_json: str, prompt: str):
     try:
         history = json.loads(history_json) if history_json else []
     except json.JSONDecodeError:
         history = []
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
     for turn in history:
         role = turn.get("role")
         content = turn.get("content", "")
@@ -280,7 +299,7 @@ def _gpu_stream(input_ids, max_new_tokens, max_iters, full_denoise, canvas_ids):
 # --------------------------------------------------------------------------- #
 # Server
 # --------------------------------------------------------------------------- #
-app = Server(title="Gemma Diffusion Website Builder")
 @app.api(name="generate", concurrency_limit=1, time_limit=600, stream_every=0.05)
@@ -292,23 +311,24 @@ def generate(
     full_denoise: bool = False,
     anim_delay: float = 0.0,
     warm_start: bool = True,
 ) -> str:
-    """Stream the diffusion generation as JSON frames (one per denoising step).
-    The model writes a self-contained HTML document; the frontend renders it live.
-    """
     prompt = (prompt or "").strip()
     if not prompt:
         yield json.dumps({"kind": "error", "message": "Empty prompt."})
         return
-    messages = build_messages(history_json, prompt)
     max_iters = max(1, min(int(max_iters), MAX_ITERS_CAP))
     # Tweak warm-start: seed the diffusion's first canvas with the previous page's own
     # tokens (native `canvas_ids` API) so the model edits the existing page in place.
     is_tweak = bool(last_assistant_html(history_json))
-    canvas_ids = warm_canvas_from_cache() if (warm_start and is_tweak) else None
     warming = canvas_ids is not None
     input_ids = tokenizer.apply_chat_template(
@@ -336,16 +356,18 @@ def generate(
                 "canvas": CANVAS_LEN,
                 "max_iters": max_iters,
                 "warming": warming,
             }
         )
         if anim_delay and kind == "draft":
             _time.sleep(float(anim_delay))
-    final_source = extract_html(last_text)
-    # Cache the *cleaned* output so the next tweak warm-starts from a valid header.
     if final_source.strip():
-        model._last_clean_html = final_source
-    yield json.dumps({"kind": "done", "source": final_source})
 @app.get("/", response_class=HTMLResponse)
@@ -362,4 +384,4 @@ if __name__ == "__main__":
         server_name=os.environ.get("GDIFF_HOST", "0.0.0.0"),
         server_port=int(os.environ.get("GDIFF_PORT", "7860")),
         show_error=True,
-    )

+# app.py
 """
 Gemma Diffusion — live website builder (gradio.Server backend + custom frontend).
 # ZeroGPU: the 26B checkpoint (~49 GB bf16) needs the full backing card.
 GPU_SIZE = os.environ.get("GDIFF_GPU_SIZE", "xlarge")
+SYSTEM_PROMPT_HTML = (
     "You are an expert front-end web developer with great visual taste. When asked to "
     "build or change a web page, respond with a SINGLE, complete, self-contained HTML5 "
     "document. Put all CSS in a <style> tag and any JavaScript in a <script> tag inside "
     "<!DOCTYPE html>."
 )
+SYSTEM_PROMPT_PYTHON = (
+    "You are an expert Python developer. When asked to write or change a Python program, "
+    "respond with a SINGLE, complete, self-contained Python script. Do not use external "
+    "packages that require installation unless explicitly requested; prefer the standard "
+    "library. When asked to modify an existing script, return the FULL updated script "
+    "with the change applied. Do not include explanations or markdown code fences — "
+    "output only raw Python code."
+)
+def system_prompt_for(mode: str) -> str:
+    return SYSTEM_PROMPT_PYTHON if (mode or "").lower() == "python" else SYSTEM_PROMPT_HTML
 _MARKER_RE = re.compile(
     r"<\|?(?:channel|turn|think|image|audio|video|tool(?:_call|_response)?)\|?>"
 )
+_FENCE_RE = re.compile(r"```(?:html|python)?\s*(.*?)\s*```", re.DOTALL)
 # --------------------------------------------------------------------------- #
 PAD_ID = tokenizer.pad_token_id or 0
 print(f"[gdiff] Model ready | canvas_length={CANVAS_LEN}", flush=True)
+# Cache of the last *cleaned* page/code so a follow-up tweak can warm-start in place.
 model._last_clean_html = None
+model._last_clean_python = None
 # --------------------------------------------------------------------------- #
 # Helpers (CPU-only; safe to run in the gradio.Server main process)
 # --------------------------------------------------------------------------- #
+def warm_canvas_from_cache(mode: str = "webpage"):
+    """Starting canvas (first block) built from the previous *cleaned* page/code.
     Returns a CPU tensor (it is pickled across the ZeroGPU process boundary and moved
+    to CUDA inside the GPU worker). We re-tokenize the cleaned HTML/python rather than
+    reuse raw output tokens so a mangled header can't compound across tweaks.
     """
+    attr = "_last_clean_python" if (mode or "").lower() == "python" else "_last_clean_html"
+    src = getattr(model, attr, None)
+    if not src:
         return None
+    ids = tokenizer(src, add_special_tokens=False).input_ids[:CANVAS_LEN]
     if not ids:
         return None
     if len(ids) < CANVAS_LEN:
 def extract_html(text: str) -> str:
+    """Pull a usable HTML document out of the (possibly mangled) model output."""
     text = clean_text(text)
     fenced = _FENCE_RE.search(text)
     if fenced:
     return text.strip()
+def extract_python(text: str) -> str:
+    """Pull usable Python code out of the (possibly mangled) model output."""
+    text = clean_text(text)
+    fenced = _FENCE_RE.search(text)
+    if fenced:
+        text = fenced.group(1)
+    return text.strip()
 class QueueDiffusionStreamer(BaseStreamer):
     def __init__(self, tok, q: "queue_lib.Queue"):
         self.tok = tok
         self.q.put(("end", self._decode(self.confirmed_ids), self.block, self.step))
+def build_messages(history_json: str, prompt: str, mode: str = "webpage"):
     try:
         history = json.loads(history_json) if history_json else []
     except json.JSONDecodeError:
         history = []
+    messages = [{"role": "system", "content": system_prompt_for(mode)}]
     for turn in history:
         role = turn.get("role")
         content = turn.get("content", "")
 # --------------------------------------------------------------------------- #
 # Server
 # --------------------------------------------------------------------------- #
+app = Server(title="Gemma Diffusion Builder")
 @app.api(name="generate", concurrency_limit=1, time_limit=600, stream_every=0.05)
     full_denoise: bool = False,
     anim_delay: float = 0.0,
     warm_start: bool = True,
+    mode: str = "webpage",
 ) -> str:
+    """Stream the diffusion generation as JSON frames (one per denoising step)."""
     prompt = (prompt or "").strip()
     if not prompt:
         yield json.dumps({"kind": "error", "message": "Empty prompt."})
         return
+    mode = (mode or "webpage").lower()
+    is_python = mode == "python"
+    messages = build_messages(history_json, prompt, mode)
     max_iters = max(1, min(int(max_iters), MAX_ITERS_CAP))
     # Tweak warm-start: seed the diffusion's first canvas with the previous page's own
     # tokens (native `canvas_ids` API) so the model edits the existing page in place.
     is_tweak = bool(last_assistant_html(history_json))
+    canvas_ids = warm_canvas_from_cache(mode) if (warm_start and is_tweak) else None
     warming = canvas_ids is not None
     input_ids = tokenizer.apply_chat_template(
                 "canvas": CANVAS_LEN,
                 "max_iters": max_iters,
                 "warming": warming,
+                "mode": mode,
             }
         )
         if anim_delay and kind == "draft":
             _time.sleep(float(anim_delay))
+    final_source = extract_python(last_text) if is_python else extract_html(last_text)
+    # Cache the *cleaned* output so the next tweak warm-starts from a valid artifact.
+    cache_attr = "_last_clean_python" if is_python else "_last_clean_html"
     if final_source.strip():
+        setattr(model, cache_attr, final_source)
+    yield json.dumps({"kind": "done", "source": final_source, "mode": mode})
 @app.get("/", response_class=HTMLResponse)
         server_name=os.environ.get("GDIFF_HOST", "0.0.0.0"),
         server_port=int(os.environ.get("GDIFF_PORT", "7860")),
         show_error=True,
+    )