padding patch

assets/index-Bf-HmMxp.js

@@ -4013,7 +4013,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let n_idx = h * W + w;
   y[c * N + n_idx] = x[n_idx * C + c];
 }
-`,d=i.empty(o,[t,n,r],`t-nc-chw`);return await i.runProgram({name:`t_nc_chw`,source:u,cacheKey:l,bindings:[{tensor:e,type:`read-only-storage`},{tensor:d,type:`storage`}],workgroups:[Math.ceil(n/8),Math.ceil(r/8),t]}),d}};function mi(e){if(e.byteLength%4==0)return e;let t=new Uint16Array(e.length+1);return t.set(e),t}function hi({rt:e,weightF16:t,inC:n,outC:r}){let i=bi(),a=[];for(let o=0;o<2;++o){let s=o===0?[[0],[1,2]]:[[0,1],[2]];for(let o=0;o<2;++o){let c=o===0?[[0],[1,2]]:[[0,1],[2]],l=new Uint16Array(r*n*4);for(let e=0;e<r;++e)for(let r=0;r<n;++r)for(let a=0;a<2;++a)for(let o=0;o<2;++o){let u=0;for(let l of s[a])for(let a of c[o])u+=i[t[e*n*9+r*9+l*3+a]];l[e*n*4+r*4+a*2+o]=It(u)}let u=mi(l);a.push(e.tensorFromTypedArray(`float16`,[u.length],u))}}return a}function gi({rt:e,weightF16:t,inC:n,outC:r}){let i=bi(),a=[];for(let o=0;o<2;++o){let s=o===0?[[0],[1,2]]:[[0,1],[2]];for(let o=0;o<2;++o){let c=o===0?[[0],[1,2]]:[[0,1],[2]],l=new Uint16Array(n*4*r);for(let e=0;e<n;++e)for(let a=0;a<2;++a)for(let o=0;o<2;++o){let u=a*2+o,d=(e*4+u)*r;for(let u=0;u<r;++u){let r=0;for(let l of s[a])for(let a of c[o])r+=i[t[u*n*9+e*9+l*3+a]];l[d+u]=It(r)}}let u=mi(l);a.push(e.tensorFromTypedArray(`float16`,[u.length],u))}}return a}function _i({rt:e,weightF16:t,inC:n,outC:r}){let i=new Uint16Array(r*9*n);for(let e=0;e<r;++e){let r=e*n*9;for(let e=0;e<9;++e){let a=r+e*n;for(let o=0;o<n;++o)i[a+o]=t[r+o*9+e]}}return e.tensorFromTypedArray(`float16`,[i.length],i)}async function vi({rt:e,weightF16:t,inC:n,outC:r}){let i=e.tensorFromTypedArray(`float16`,[t.length],t),a=e.empty(`float16`,[n*16,r],`conv-wino-f2x2-weight`);return await e.runProgram({name:`conv2d_winograd_weight_transform`,source:ei({inC:n,outC:r}),cacheKey:`conv2d_winograd_weight_transform_${n}_${r}`,bindings:[{tensor:i,type:`read-only-storage`},{tensor:a,type:`storage`}],workgroups:[Math.ceil(r/16),Math.ceil(n/16),1]}),e.clearBindGroupCache?.(),e.host.device.queue.onSubmittedWorkDone().then(()=>i.destroy?.()).catch(()=>{}),a}var yi=null;function bi(){if(yi)return yi;let e=new Float32Array(65536);for(let t=0;t<e.length;++t)e[t]=Ft(t);return yi=e,e}function xi(e){let t=new Float32Array(e*4);for(let n=0;n<e;++n)t[n*4+3]=n;return t}function Si(e,t){let n=new Float32Array(e*t*4);for(let r=0;r<e;++r)for(let e=0;e<t;++e){let i=(r*t+e)*4;n[i+1]=r,n[i+2]=e}return n}var Ci=class{constructor({rt:e,snapshotDir:t,tokenizer:n,textEncoder:r,transformer:i,vae:a,vaeConfig:o,schedulerConfig:s,bnStats:c}){this.rt=e,this.snapshotDir=t,this.tokenizer=n,this.textEncoder=r,this.transformer=i,this.vae=a,this.vaeConfig=o??a?.config??null,this.schedulerConfig=s,this.bnStats=c,this.destroyed=!1}async ensureVae(){if(this.vae&&this.bnStats)return this.vae;throw Error(`VAE was not loaded; construct the pipeline without skipVae to decode images`)}async generate(e){let t=Vt();try{return await this._generate(e,t)}finally{this.rt.clearTransientCaches?.(),this.rt.clearReadbackPool?.(),t.destroy(),this.rt.clearTransientCaches?.()}}async _generate({prompt:e,height:t=1024,width:n=1024,numInferenceSteps:r=4,seed:i=0,log:a=null,callbackOnStepEnd:o=null,encoderHiddenStatesT:s=null},c=null){if(t%16!=0||n%16!=0)throw Error(`height and width must be divisible by 16`);let l=this.rt,u=this.schedulerConfig,d=this.transformer.config;d.num_attention_heads*d.attention_head_dim;let f=s,p,m;if(f){if(f.runtime!==l)throw Error(`encoderHiddenStatesT belongs to a different runtime`);if(f.shape.length!==2)throw Error(`encoderHiddenStatesT must have shape [seq, stackDim]`);[p,m]=f.shape,a?.(`text encode cache`)}else ({hiddenStackT:f,seq:p,stackDim:m}=await this.encodePrompt(e,{log:a,scope:c}));if(m!==d.joint_attention_dim)throw Error(`text stackDim ${m} != joint_attention_dim ${d.joint_attention_dim}`);a?.(`scheduler`);let h=dr(t/16*(n/16),r),g=new pr(u);g.setTimesteps({numInferenceSteps:r,mu:h});let _=t/8,v=n/8,y=this.vaeConfig?.latent_channels??32,b=_/2,x=v/2,S=y*4,C=Ii(Fi(i,S*b*x),S,b,x),w=c?.track(l.tensorFromTypedArray(`float32`,[b*x,S],C))??l.tensorFromTypedArray(`float32`,[b*x,S],C),T=Si(b,x),E=xi(p);for(let e=0;e<r;++e){let t=Vt(),n=Ht(l,t);try{let i=g.timesteps[e]/1e3;a?.(`step ${e}/${r} t=${i.toFixed(4)}`);let s=await this.transformer.forward({hiddenStatesT:w,encoderHiddenStatesT:f,timestep:i,imgIds:T,txtIds:E,scope:t}),c=g.stepDelta(e);await Ln(n,{xT:w,yT:s,count:b*x*128,alpha:c}),o&&await o(this,e,g.timesteps[e],{latents:w})}finally{l.clearBindGroupCache?.(),t.destroy()}}await this.ensureVae(),a?.(`unpack + BN-denorm`);let D=await Rn(c?Ht(l,c):l,{packedT:w,meanT:this.bnStats.running_meanT,stdT:this.bnStats.running_stdT,outputDtype:l.caps().f16?`float16`:`float32`,latentC:y,latentH:_,latentW:v});a?.(`vae decode`);let{image:O,H:k,W:A}=await this.vae.decode(D,_,v,{scope:c}),j=await Ei(l,O);a?.(`to RGB`);let M=new Uint8Array(k*A*3);for(let e=0;e<3;e++)for(let t=0;t<k;t++)for(let n=0;n<A;n++){let r=(j[e*k*A+t*A+n]+1)*127.5;M[(t*A+n)*3+e]=Math.min(255,Math.max(0,Math.round(r)))}return a?.(`png encode`),br(A,k,M)}destroy(){this.destroyed||(this.destroyed=!0,this.rt.clearTransientCaches?.(),this.textEncoder?.destroy?.(),this.transformer?.destroy?.(),this.vae?.destroy?.(),this.tokenizer=null,this.textEncoder=null,this.transformer=null,this.vae=null,this.bnStats=null,this.rt.clearTransientCaches?.())}async encodePrompt(e,{log:t=null,scope:n=null}={}){if(!this.tokenizer||!this.textEncoder)throw Error(`Text encoder was not loaded; provide encoderHiddenStatesT to generate()`);t?.(`tokenize`);let r=Pi(e),i=(await this.tokenizer.encode(r)).ids.slice(0,512),a=new Uint32Array(Math.max(1,i.length));for(let e=0;e<i.length;++e)a[e]=i[e];return t?.(`tokens: ${i.length}`),t?.(`text encode`),this.textEncoder.encode(a,{scope:n})}},wi=class e extends Ci{static async fromSnapshot(t,n,r={}){return Ti(e,t,n,r,{readJsonResource:Di,readJsonResourceOptional:Oi,openSafeTensorsResource:ki})}};async function Ti(e,t,n,{onProgress:r=null,fetch:i=null,cacheStorage:a=null,cacheName:o=null,cache:s=void 0,force:c=!1,signal:l=null,skipTextEncoder:u=!1,skipVae:d=!1,requireRangeRequests:f=!0}={},p){let m={fetch:i,cacheStorage:a,cacheName:o,cache:s,force:c,signal:l,requireRangeRequests:f},h=null,g=null,_=null,v=e=>{r&&r({component:e,phase:`open`})},y=e=>t=>{r&&r({component:e,phase:`download`,...t})};try{r&&r({phase:`init`});let i=await p.readJsonResource(n,`scheduler/scheduler_config.json`,m),a=new At(await p.readJsonResource(n,`tokenizer/tokenizer.json`,m),await p.readJsonResource(n,`tokenizer/tokenizer_config.json`,m));if(!u){v(`text_encoder`);let e=await p.readJsonResource(n,`text_encoder-mlx-4bit/config.json`,m),r=await p.openSafeTensorsResource(n,`text_encoder-mlx-4bit/model.safetensors`,m);try{h=await cr.fromMlxSafeTensors({rt:t,config:e,safeTensors:r,onProgress:y(`text_encoder`),signal:l})}finally{await r.close()}}v(`transformer`);let o=await p.readJsonResource(n,`transformer-packed-mflux/config.json`,m),s=p.readJsonResourceOptional?await p.readJsonResourceOptional(n,`transformer-packed-mflux/quantization_config.json`,m):null,c=s?{...o,quantization_config:s}:o,f=await p.openSafeTensorsResource(n,`transformer-packed-mflux/diffusion_pytorch_model.safetensors`,m);try{g=await Nr.fromMlxSafeTensors({rt:t,config:c,safeTensors:f,onProgress:y(`transformer`),signal:l})}finally{await f.close()}let b=await p.readJsonResource(n,`vae/config.json`,m),x=null;if(!d){v(`vae`);let e=await p.openSafeTensorsResource(n,`vae/diffusion_pytorch_model.safetensors`,m);try{_=await pi.fromBf16SafeTensors({rt:t,config:b,safeTensors:e,onProgress:y(`vae`),signal:l}),x=_.w.bn}finally{await e.close()}}return new e({rt:t,snapshotDir:n,tokenizer:a,textEncoder:h,transformer:g,vae:_,vaeConfig:b,schedulerConfig:i,bnStats:x})}catch(e){throw h?.destroy?.(),g?.destroy?.(),_?.destroy?.(),t.clearTransientCaches?.(),e}}async function Ei(e,t){let n=await e.readTensor(t);return t.dtype===`float16`?Float32Array.from(n,Ft):n}async function Di(e,t,n){return JSON.parse(await ji(Ai(e,t),n))}async function Oi(e,t,n){try{return await Di(e,t,n)}catch{return null}}async function ki(e,t,n={}){return A(Ai(e,t),n)}function Ai(e,t){let n=e instanceof URL?e.toString():String(e);return new URL(t,Ni(n),globalThis.location?.href).toString()}async function ji(e,t={}){let n=t.fetch??globalThis.fetch;if(typeof n!=`function`)throw Error(`No fetch implementation available`);let r=await Mi(t);if(r&&!t.force){let t=await r.match(e);if(t)return t.text()}let i=await n(e,{signal:t.signal});if(!i.ok)throw Error(`GET ${e} failed: ${i.status} ${i.statusText}`);if(r)try{await r.put(e,i.clone())}catch(e){typeof console<`u`&&console.warn(`resource cache write failed: ${e.message}`)}return i.text()}async function Mi(e={}){if(e.cache===!1)return null;let t=e.cacheStorage??globalThis.caches;return t?.open?t.open(e.cacheName??`bonsai-pipeline-v1`):null}function Ni(e){return e.endsWith(`/`)?e:`${e}/`}function Pi(e){return`<|im_start|>user
+`,d=i.empty(o,[t,n,r],`t-nc-chw`);return await i.runProgram({name:`t_nc_chw`,source:u,cacheKey:l,bindings:[{tensor:e,type:`read-only-storage`},{tensor:d,type:`storage`}],workgroups:[Math.ceil(n/8),Math.ceil(r/8),t]}),d}};function mi(e){if(e.byteLength%4==0)return e;let t=new Uint16Array(e.length+1);return t.set(e),t}function hi({rt:e,weightF16:t,inC:n,outC:r}){let i=bi(),a=[];for(let o=0;o<2;++o){let s=o===0?[[0],[1,2]]:[[0,1],[2]];for(let o=0;o<2;++o){let c=o===0?[[0],[1,2]]:[[0,1],[2]],l=new Uint16Array(r*n*4);for(let e=0;e<r;++e)for(let r=0;r<n;++r)for(let a=0;a<2;++a)for(let o=0;o<2;++o){let u=0;for(let l of s[a])for(let a of c[o])u+=i[t[e*n*9+r*9+l*3+a]];l[e*n*4+r*4+a*2+o]=It(u)}let u=mi(l);a.push(e.tensorFromTypedArray(`float16`,[u.length],u))}}return a}function gi({rt:e,weightF16:t,inC:n,outC:r}){let i=bi(),a=[];for(let o=0;o<2;++o){let s=o===0?[[0],[1,2]]:[[0,1],[2]];for(let o=0;o<2;++o){let c=o===0?[[0],[1,2]]:[[0,1],[2]],l=new Uint16Array(n*4*r);for(let e=0;e<n;++e)for(let a=0;a<2;++a)for(let o=0;o<2;++o){let u=a*2+o,d=(e*4+u)*r;for(let u=0;u<r;++u){let r=0;for(let l of s[a])for(let a of c[o])r+=i[t[u*n*9+e*9+l*3+a]];l[d+u]=It(r)}}let u=mi(l);a.push(e.tensorFromTypedArray(`float16`,[u.length],u))}}return a}function _i({rt:e,weightF16:t,inC:n,outC:r}){let i=new Uint16Array(r*9*n);for(let e=0;e<r;++e){let r=e*n*9;for(let e=0;e<9;++e){let a=r+e*n;for(let o=0;o<n;++o)i[a+o]=t[r+o*9+e]}}return e.tensorFromTypedArray(`float16`,[i.length],i)}async function vi({rt:e,weightF16:t,inC:n,outC:r}){let i=e.tensorFromTypedArray(`float16`,[t.length],t),a=e.empty(`float16`,[n*16,r],`conv-wino-f2x2-weight`);return await e.runProgram({name:`conv2d_winograd_weight_transform`,source:ei({inC:n,outC:r}),cacheKey:`conv2d_winograd_weight_transform_${n}_${r}`,bindings:[{tensor:i,type:`read-only-storage`},{tensor:a,type:`storage`}],workgroups:[Math.ceil(r/16),Math.ceil(n/16),1]}),e.clearBindGroupCache?.(),e.host.device.queue.onSubmittedWorkDone().then(()=>i.destroy?.()).catch(()=>{}),a}var yi=null;function bi(){if(yi)return yi;let e=new Float32Array(65536);for(let t=0;t<e.length;++t)e[t]=Ft(t);return yi=e,e}function xi(e){let t=new Float32Array(e*4);for(let n=0;n<e;++n)t[n*4+3]=n;return t}function Si(e,t){let n=new Float32Array(e*t*4);for(let r=0;r<e;++r)for(let e=0;e<t;++e){let i=(r*t+e)*4;n[i+1]=r,n[i+2]=e}return n}var Ci=class{constructor({rt:e,snapshotDir:t,tokenizer:n,textEncoder:r,transformer:i,vae:a,vaeConfig:o,schedulerConfig:s,bnStats:c}){this.rt=e,this.snapshotDir=t,this.tokenizer=n,this.textEncoder=r,this.transformer=i,this.vae=a,this.vaeConfig=o??a?.config??null,this.schedulerConfig=s,this.bnStats=c,this.destroyed=!1}async ensureVae(){if(this.vae&&this.bnStats)return this.vae;throw Error(`VAE was not loaded; construct the pipeline without skipVae to decode images`)}async generate(e){let t=Vt();try{return await this._generate(e,t)}finally{this.rt.clearTransientCaches?.(),this.rt.clearReadbackPool?.(),t.destroy(),this.rt.clearTransientCaches?.()}}async _generate({prompt:e,height:t=1024,width:n=1024,numInferenceSteps:r=4,seed:i=0,log:a=null,callbackOnStepEnd:o=null,encoderHiddenStatesT:s=null},c=null){if(t%16!=0||n%16!=0)throw Error(`height and width must be divisible by 16`);let l=this.rt,u=this.schedulerConfig,d=this.transformer.config;d.num_attention_heads*d.attention_head_dim;let f=s,p,m;if(f){if(f.runtime!==l)throw Error(`encoderHiddenStatesT belongs to a different runtime`);if(f.shape.length!==2)throw Error(`encoderHiddenStatesT must have shape [seq, stackDim]`);[p,m]=f.shape,a?.(`text encode cache`)}else ({hiddenStackT:f,seq:p,stackDim:m}=await this.encodePrompt(e,{log:a,scope:c}));if(m!==d.joint_attention_dim)throw Error(`text stackDim ${m} != joint_attention_dim ${d.joint_attention_dim}`);a?.(`scheduler`);let h=dr(t/16*(n/16),r),g=new pr(u);g.setTimesteps({numInferenceSteps:r,mu:h});let _=t/8,v=n/8,y=this.vaeConfig?.latent_channels??32,b=_/2,x=v/2,S=y*4,C=Ii(Fi(i,S*b*x),S,b,x),w=c?.track(l.tensorFromTypedArray(`float32`,[b*x,S],C))??l.tensorFromTypedArray(`float32`,[b*x,S],C),T=Si(b,x),E=xi(p);for(let e=0;e<r;++e){let t=Vt(),n=Ht(l,t);try{let i=g.timesteps[e]/1e3;a?.(`step ${e}/${r} t=${i.toFixed(4)}`);let s=await this.transformer.forward({hiddenStatesT:w,encoderHiddenStatesT:f,timestep:i,imgIds:T,txtIds:E,scope:t}),c=g.stepDelta(e);await Ln(n,{xT:w,yT:s,count:b*x*128,alpha:c}),o&&await o(this,e,g.timesteps[e],{latents:w})}finally{l.clearBindGroupCache?.(),t.destroy()}}await this.ensureVae(),a?.(`unpack + BN-denorm`);let D=await Rn(c?Ht(l,c):l,{packedT:w,meanT:this.bnStats.running_meanT,stdT:this.bnStats.running_stdT,outputDtype:l.caps().f16?`float16`:`float32`,latentC:y,latentH:_,latentW:v});a?.(`vae decode`);let{image:O,H:k,W:A}=await this.vae.decode(D,_,v,{scope:c}),j=await Ei(l,O);a?.(`to RGB`);let M=new Uint8Array(k*A*3);for(let e=0;e<3;e++)for(let t=0;t<k;t++)for(let n=0;n<A;n++){let r=(j[e*k*A+t*A+n]+1)*127.5;M[(t*A+n)*3+e]=Math.min(255,Math.max(0,Math.round(r)))}return a?.(`png encode`),br(A,k,M)}destroy(){this.destroyed||(this.destroyed=!0,this.rt.clearTransientCaches?.(),this.textEncoder?.destroy?.(),this.transformer?.destroy?.(),this.vae?.destroy?.(),this.tokenizer=null,this.textEncoder=null,this.transformer=null,this.vae=null,this.bnStats=null,this.rt.clearTransientCaches?.())}async encodePrompt(e,{log:t=null,scope:n=null}={}){if(!this.tokenizer||!this.textEncoder)throw Error(`Text encoder was not loaded; provide encoderHiddenStatesT to generate()`);t?.(`tokenize`);let r=Pi(e),i=(await this.tokenizer.encode(r)).ids.slice(0,512),a=new Uint32Array(512),o=this.tokenizer.config?.pad_token,s=typeof o==`string`?o:o?.content,c=s&&(this.tokenizer.token_to_id?.(s)??this.tokenizer.tokenToId?.(s));a.fill(Number.isInteger(c)&&c>=0?c:0);for(let e=0;e<i.length;++e)a[e]=i[e];return t?.(`tokens: ${i.length}`),t?.(`text encode`),this.textEncoder.encode(a,{scope:n})}},wi=class e extends Ci{static async fromSnapshot(t,n,r={}){return Ti(e,t,n,r,{readJsonResource:Di,readJsonResourceOptional:Oi,openSafeTensorsResource:ki})}};async function Ti(e,t,n,{onProgress:r=null,fetch:i=null,cacheStorage:a=null,cacheName:o=null,cache:s=void 0,force:c=!1,signal:l=null,skipTextEncoder:u=!1,skipVae:d=!1,requireRangeRequests:f=!0}={},p){let m={fetch:i,cacheStorage:a,cacheName:o,cache:s,force:c,signal:l,requireRangeRequests:f},h=null,g=null,_=null,v=e=>{r&&r({component:e,phase:`open`})},y=e=>t=>{r&&r({component:e,phase:`download`,...t})};try{r&&r({phase:`init`});let i=await p.readJsonResource(n,`scheduler/scheduler_config.json`,m),a=new At(await p.readJsonResource(n,`tokenizer/tokenizer.json`,m),await p.readJsonResource(n,`tokenizer/tokenizer_config.json`,m));if(!u){v(`text_encoder`);let e=await p.readJsonResource(n,`text_encoder-mlx-4bit/config.json`,m),r=await p.openSafeTensorsResource(n,`text_encoder-mlx-4bit/model.safetensors`,m);try{h=await cr.fromMlxSafeTensors({rt:t,config:e,safeTensors:r,onProgress:y(`text_encoder`),signal:l})}finally{await r.close()}}v(`transformer`);let o=await p.readJsonResource(n,`transformer-packed-mflux/config.json`,m),s=p.readJsonResourceOptional?await p.readJsonResourceOptional(n,`transformer-packed-mflux/quantization_config.json`,m):null,c=s?{...o,quantization_config:s}:o,f=await p.openSafeTensorsResource(n,`transformer-packed-mflux/diffusion_pytorch_model.safetensors`,m);try{g=await Nr.fromMlxSafeTensors({rt:t,config:c,safeTensors:f,onProgress:y(`transformer`),signal:l})}finally{await f.close()}let b=await p.readJsonResource(n,`vae/config.json`,m),x=null;if(!d){v(`vae`);let e=await p.openSafeTensorsResource(n,`vae/diffusion_pytorch_model.safetensors`,m);try{_=await pi.fromBf16SafeTensors({rt:t,config:b,safeTensors:e,onProgress:y(`vae`),signal:l}),x=_.w.bn}finally{await e.close()}}return new e({rt:t,snapshotDir:n,tokenizer:a,textEncoder:h,transformer:g,vae:_,vaeConfig:b,schedulerConfig:i,bnStats:x})}catch(e){throw h?.destroy?.(),g?.destroy?.(),_?.destroy?.(),t.clearTransientCaches?.(),e}}async function Ei(e,t){let n=await e.readTensor(t);return t.dtype===`float16`?Float32Array.from(n,Ft):n}async function Di(e,t,n){return JSON.parse(await ji(Ai(e,t),n))}async function Oi(e,t,n){try{return await Di(e,t,n)}catch{return null}}async function ki(e,t,n={}){return A(Ai(e,t),n)}function Ai(e,t){let n=e instanceof URL?e.toString():String(e);return new URL(t,Ni(n),globalThis.location?.href).toString()}async function ji(e,t={}){let n=t.fetch??globalThis.fetch;if(typeof n!=`function`)throw Error(`No fetch implementation available`);let r=await Mi(t);if(r&&!t.force){let t=await r.match(e);if(t)return t.text()}let i=await n(e,{signal:t.signal});if(!i.ok)throw Error(`GET ${e} failed: ${i.status} ${i.statusText}`);if(r)try{await r.put(e,i.clone())}catch(e){typeof console<`u`&&console.warn(`resource cache write failed: ${e.message}`)}return i.text()}async function Mi(e={}){if(e.cache===!1)return null;let t=e.cacheStorage??globalThis.caches;return t?.open?t.open(e.cacheName??`bonsai-pipeline-v1`):null}function Ni(e){return e.endsWith(`/`)?e:`${e}/`}function Pi(e){return`<|im_start|>user
 ${e}<|im_end|>
 <|im_start|>assistant
 <think>