Make example switching reuse preloaded images and sessions instead of re-encoding and downloading

.gitattributes

@@ -35,3 +35,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 examples/chinese.png filter=lfs diff=lfs merge=lfs -text
 examples/english.png filter=lfs diff=lfs merge=lfs -text
+examples/chinese-preview.webp filter=lfs diff=lfs merge=lfs -text
+examples/english-preview.webp filter=lfs diff=lfs merge=lfs -text

app.py

@@ -18,6 +18,7 @@ import urllib.request
 import uuid
 import zipfile
 from concurrent.futures import ThreadPoolExecutor
+from functools import lru_cache
 from pathlib import Path
 
 from fastapi import FastAPI, File, HTTPException, UploadFile
@@ -53,7 +54,23 @@ if os.environ.get("PIXTEXT_GRADIO") == "1":
 else:
     app = FastAPI(title="PixText Studio", docs_url=None, redoc_url=None)
 app.mount("/static", StaticFiles(directory=ROOT / "static"), name="static")
-app.mount("/examples", StaticFiles(directory=ROOT / "examples"), name="examples")
+
+class ExampleFiles(StaticFiles):
+    async def get_response(self, path, scope):
+        response = await super().get_response(path, scope)
+        if path.endswith("-preview.webp") and response.status_code in (200, 304):
+            response.headers["Cache-Control"] = "public, max-age=31536000, immutable"
+        return response
+
+
+app.mount("/examples", ExampleFiles(directory=ROOT / "examples"), name="examples")
+
+
+@lru_cache(maxsize=2)
+def example_preview(example: str):
+    path = ROOT / "examples" / f"{example}-preview.webp"
+    version = hashlib.sha256(path.read_bytes()).hexdigest()[:12]
+    return f"/examples/{path.name}?v={version}"
 
 
 class ExampleRequest(BaseModel):
@@ -109,14 +126,19 @@ def make_session(image: Image.Image, example: str | None = None):
         sid = uuid.uuid4().hex
         directory = RUNTIME / sid
         directory.mkdir()
-        image.save(directory / "source.png")
-        image.save(directory / "preview.webp", quality=94, method=4)
+        if example:
+            shutil.copyfile(ROOT / "examples" / f"{example}.png", directory / "source.png")
+            shutil.copyfile(ROOT / "examples" / f"{example}-preview.webp", directory / "preview.webp")
+        else:
+            image.save(directory / "source.png")
+            image.save(directory / "preview.webp", quality=94, method=4)
         SESSIONS[sid] = {
             "directory": directory, "size": image.size, "example": example,
             "touched": time.monotonic(), "job": None,
         }
     return {"id": sid, "width": image.width, "height": image.height,
-            "image": f"/api/sessions/{sid}/preview.webp", "example": example}
+            "image": example_preview(example) if example else f"/api/sessions/{sid}/preview.webp",
+            "example": example}
 
 
 def update_job(jid: str, **values):
@@ -320,7 +342,7 @@ def example_session(request: ExampleRequest):
     if request.example not in EXAMPLES:
         raise HTTPException(404, "Example not found")
     with Image.open(ROOT / "examples" / f"{request.example}.png") as image:
-        session = make_session(image.convert("RGB"), request.example)
+        session = make_session(image, request.example)
     return {**session, **example_preset(request.example)}
 
 

static/app.js

@@ -6,6 +6,22 @@ let gesture = null;
 const clone = (x) => structuredClone(x);
 const escape = (s) => String(s).replace(/[&<>"']/g, c => ({'&':'&amp;','<':'&lt;','>':'&gt;','"':'&quot;',"'":'&#39;'}[c]));
 const endpoint = (name) => `/api/sessions/${state.session.id}/${name}`;
+const exampleCache = new Map();
+
+function prepareExample(example) {
+  const cached = exampleCache.get(example);
+  if(cached && cached.expires > Date.now()) return cached.promise;
+  const entry = {expires:Date.now()+45*60*1000};
+  entry.promise = (async()=>{
+    const session = await api('/api/session',{example});
+    const image = new Image(); image.src=session.image;
+    await image.decode();
+    return {session,image};
+  })();
+  exampleCache.set(example,entry);
+  entry.promise.catch(()=>{if(exampleCache.get(example)===entry)exampleCache.delete(example);});
+  return entry.promise;
+}
 
 async function api(url, body, isForm=false) {
   const response = await fetch(url, body === undefined ? {} : {method:'POST', headers:isForm?{}:{'Content-Type':'application/json'}, body:isForm?body:JSON.stringify(body)});
@@ -28,13 +44,13 @@ function changed() {
   updateControls();
 }
 function selected() {return state.regions.find(r=>r.id===state.selected);}
-function setBusy(busy) {
+function setBusy(busy, message='Waiting for the model') {
   state.busy=busy;
   document.querySelectorAll('[data-example],[data-model],#uploadButton,#exampleTargets,#resetButton,#drawTool,#undoButton,#targetText,#enableRegion,#deleteRegion,.coords input,.settings input,.region-row input').forEach(e=>e.disabled=busy);
   $('regionCanvas').style.pointerEvents=busy?'none':'';
   $('jobProgress').hidden=!busy;
   $('jobProgress').classList.toggle('waiting',busy);
-  if(busy){$('jobMessage').textContent='Waiting for the model';$('jobPercent').textContent='';$('progressFill').style.width='0%';}
+  if(busy){$('jobMessage').textContent=message;$('jobPercent').textContent='';$('progressFill').style.width='0%';}
   updateControls();
 }
 function updateControls() {
@@ -63,7 +79,8 @@ function applySettings(settings) {
     if(settings[key]!=null)$(id).value=settings[key];
   }
 }
-async function loadSession(session) {
+async function loadSession(session, source=null) {
+  if(!source){source=new Image();source.src=session.image;await source.decode();}
   state.session=session;state.regions=clone(session.regions||[]);state.selected=state.regions[0]?.id||null;state.history=[];state.result=null;state.stale=false;state.tool='select';
   if(session.settings)applySettings(session.settings);
   $('regionResults').hidden=true;$('compareView').disabled=true;
@@ -72,14 +89,18 @@ async function loadSession(session) {
   const image=$('sourceImage'); image.setAttribute('href',session.image); image.setAttribute('width',session.width);image.setAttribute('height',session.height);
   $('imageSize').textContent=`${session.width} × ${session.height}`;
   $('wipe').style.aspectRatio=`${session.width} / ${session.height}`;
-  state.source=new Image(); state.source.src=session.image;
-  await state.source.decode();
+  state.source=source;
   document.querySelectorAll('[data-example]').forEach(e=>e.classList.toggle('active',e.dataset.example===session.example));
   switchView('edit');setTool('select');render();notice();
 }
 async function loadExample(example) {
   if(state.busy)return;
-  try{setBusy(true);await loadSession(await api('/api/session',{example}));history.replaceState(null,'',location.pathname);}catch(e){notice(e.message);}finally{setBusy(false);}
+  try{
+    setBusy(true,'Loading example');
+    const prepared=await prepareExample(example);
+    await loadSession(clone(prepared.session),prepared.image);
+    history.replaceState(null,'',location.pathname);
+  }catch(e){notice(e.message);}finally{setBusy(false);}
 }
 async function upload(file) {
   if(!file||state.busy)return;
@@ -272,6 +293,8 @@ async function openStudio(){
     }catch(e){notice('This preview has expired. Choose an example to start a new image.');}
     finally{setBusy(false);}
   }
+  // Prepare both examples together; switching reuses decoded images and untouched presets.
+  prepareExample('chinese').catch(()=>{});
   await loadExample('english');
 }
 openStudio();

static/index.html

@@ -5,7 +5,9 @@
   <meta name="theme-color" content="#faf9f6">
   <title>PixText · Text repair studio</title>
   <link rel="stylesheet" href="/static/style.css">
-  <link rel="modulepreload" href="/static/gradio-client.js">
+  <link rel="preload" as="image" href="/examples/english-preview.webp?v=8a15b66d2fd9">
+<link rel="preload" as="image" href="/examples/chinese-preview.webp?v=dfbb3fb24330">
+<link rel="modulepreload" href="/static/gradio-client.js">
   <script defer src="/static/app.js"></script>
 </head>
 <body>

tests/test_api.py

@@ -92,3 +92,25 @@ def test_sessions_and_jobs_expire_together(client, monkeypatch):
     new_session(client)
     assert not directory.exists()
     assert client.get(f"/api/sessions/{sid}/source.png").status_code == 404
+
+
+@pytest.mark.parametrize('example', ['english', 'chinese'])
+def test_example_selection_reuses_preview_without_reencoding(client, monkeypatch, example):
+    def unexpected_encoding(*args, **kwargs):
+        raise AssertionError('Example selection is re-encoding images on the request path')
+
+    with monkeypatch.context() as patch:
+        patch.setattr(Image.Image, 'save', unexpected_encoding)
+        first = client.post('/api/session', json={'example': example}).json()
+        second = client.post('/api/session', json={'example': example}).json()
+    assert first['id'] != second['id']
+    assert first['image'] == second['image']
+    preview = client.get(first['image'])
+    assert preview.status_code == 200
+    assert 'public' in preview.headers['cache-control']
+    assert 'immutable' in preview.headers['cache-control']
+    assert client.get(first['image'], headers={'If-None-Match': preview.headers['etag']}).status_code == 304
+    assert first['image'] in client.get('/').text
+    assert Image.open(io.BytesIO(preview.content)).size == (first['width'], first['height'])
+    original = client.get(f"/api/sessions/{first['id']}/source.png")
+    assert original.headers['cache-control'] == 'private, no-store'

tests/test_examples.mjs

@@ -0,0 +1,86 @@
+import assert from 'node:assert/strict';
+import {readFileSync} from 'node:fs';
+import test from 'node:test';
+import vm from 'node:vm';
+
+const script = readFileSync(new URL('../static/app.js', import.meta.url), 'utf8');
+const presets = Object.fromEntries(['english', 'chinese'].map(name => [name,
+  JSON.parse(readFileSync(new URL(`../examples/${name}.json`, import.meta.url), 'utf8'))]));
+
+// Run the real editor against a small DOM/fetch adapter; no models or browser packages needed.
+function editor() {
+  const elements = new Map(), requests = [], decoded = [];
+  let nextId = 0, failExample = null;
+  const element = id => {
+    if(!elements.has(id)) elements.set(id, {value:'', dataset:{}, style:{setProperty(){}},
+      classList:{toggle(){},add(){},remove(){}}, setAttribute(){}, addEventListener(){},
+      getContext(){return {drawImage(){}};}, children:[], focus(){}});
+    return elements.get(id);
+  };
+  const examples = ['english','chinese'].map(example => ({...element(example),dataset:{example}}));
+  const models = ['meanflow','teacher'].map(model => ({...element(model),dataset:{model}}));
+  const context = vm.createContext({structuredClone, URLSearchParams, console, setTimeout,
+    document:{getElementById:element,addEventListener(){},querySelectorAll(selector){
+      if(selector==='[data-example]') return examples;
+      if(selector==='[data-model]') return models;
+      return [];
+    }},
+    Image:class {async decode(){decoded.push(this.src);this.complete=true;}},
+    location:{pathname:'/',hash:''}, history:{replaceState(){}},
+    fetch:async(url,options) => {
+      if(url==='/api/config') return {ok:true,json:async()=>({transport:'http'})};
+      assert.equal(url,'/api/session');
+      const {example}=JSON.parse(options.body);requests.push(example);
+      if(example===failExample){failExample=null;throw new Error('Network unavailable');}
+      return {ok:true,json:async()=>({id:`session-${++nextId}`,example,width:1024,height:1024,
+        image:`/examples/${example}-preview.webp?v=release`,
+        regions:presets[example].map((r,i)=>({...r,id:`example-${i}`,enabled:true,origin:'example'})),
+        settings:{variant:'meanflow',steps:4,seed:42,guidance:2.6,text_threshold:.5}})};
+    }});
+  vm.runInContext(script,context);
+  const run=code=>vm.runInContext(code,context);
+  return {run,requests,decoded,element,failNext(example){failExample=example;}};
+}
+
+async function ready(app) {
+  await app.run("Promise.all([prepareExample('english'),prepareExample('chinese')])");
+  await new Promise(resolve=>setImmediate(resolve));
+}
+
+test('switching and restoring example edits reuse prepared sessions and decoded images', async()=>{
+  const app=editor();await ready(app);
+  assert.equal(app.requests.length,2);
+  assert.equal(app.decoded.length,2);
+  for(const example of ['chinese','english','chinese','english']) {
+    await app.run(`loadExample('${example}')`);
+    assert.equal(app.run('state.regions.length'),presets[example].length);
+    assert.equal(app.element('guidance').value,2.6);
+    assert.equal(app.element('refineButton').disabled,false);
+    app.run("state.regions[0].text='edited'; state.regions[0].bbox[0]=0; state.regions[0].enabled=false;");
+  }
+  await app.run("loadExample('english')");
+  assert.equal(app.run('state.regions[0].text'),presets.english[0].text);
+  assert.equal(app.run('state.regions[0].bbox[0]'),presets.english[0].bbox[0]);
+  assert.equal(app.run('state.regions[0].enabled'),true);
+  assert.equal(app.requests.length,2);
+  assert.equal(app.decoded.length,2);
+});
+
+test('concurrent preparation shares one request and an expired session is renewed',async()=>{
+  const app=editor();await ready(app);
+  app.run("exampleCache.get('chinese').expires=0");
+  await app.run("Promise.all([prepareExample('chinese'),prepareExample('chinese')])");
+  assert.equal(app.requests.filter(x=>x==='chinese').length,2);
+});
+
+test('failed preparation is retryable without replacing the current editor',async()=>{
+  const app=editor();await ready(app);
+  await app.run("loadExample('english')");
+  const current=app.run('state.session.id');
+  app.run("exampleCache.get('chinese').expires=0");app.failNext('chinese');
+  await app.run("loadExample('chinese')");
+  assert.equal(app.run('state.session.id'),current);
+  assert.equal(app.run('state.busy'),false);
+  await app.run("loadExample('chinese')");
+  assert.equal(app.run('state.regions.length'),13);
+});