Rellenar: lados en potencias de 2, el resto en espejo (la FFT de la CPU solo es rápida así)

ONNX Runtime en la CPU solo hace rápida la FFT de LaMa con lados que son potencias de 2:
512 × 512 son 2,7 s y 704 × 448, casi los mismos píxeles, 10 s. El trozo se reduce lo justo
para caber en 512 (o en 1024 si pasa de 768 px) y se completa hasta la potencia de 2 con su
propio borde en espejo, sin hueco; al volver se descarta. Un trozo de 700 px pasa de 10 s
a 4 s; uno de 300 × 200, 2 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Este commit está contenido en:
Munka
2026-10-14 20:21:16 -05:00
padre eec48362e5
commit ed4c0ed529
Se han modificado 2 ficheros con 91 adiciones y 51 borrados

Ver fichero

@@ -255,17 +255,17 @@ fn rellenar(rgba: &[u8], ancho: usize, alto: usize) -> Result<Vec<u8>, String> {
let Some(a) = m.as_mut() else { let Some(a) = m.as_mut() else {
return Err("La IA no está preparada.".into()); return Err("La IA no está preparada.".into());
}; };
let (mw, mh) = relleno::tamano_modelo(ancho, alto); let p = relleno::plan(ancho, alto);
let t0 = Instant::now(); let t0 = Instant::now();
let r = a.motor.correr_forma(&[1, 4, mh, mw], relleno::entrada(rgba, ancho, alto, mw, mh)); let r = a.motor.correr_forma(&[1, 4, p.ph, p.pw], relleno::entrada(rgba, ancho, alto, p));
a.usado = Instant::now(); a.usado = Instant::now();
match r { match r {
Ok((_, planos)) if planos.len() == 3 * mw * mh => { Ok((_, planos)) if planos.len() == 3 * p.pw * p.ph => {
crate::salud::anotar("INFO", "ia", &format!("relleno de {ancho}×{alto} (a {mw}×{mh}) en {:.1} s", t0.elapsed().as_secs_f32())); crate::salud::anotar("INFO", "ia", &format!("relleno de {ancho}×{alto} (a {}×{} en {}×{}) en {:.1} s", p.sw, p.sh, p.pw, p.ph, t0.elapsed().as_secs_f32()));
let rgb = relleno::salida(&planos, mw, mh, rgba, ancho, alto); let rgb = relleno::salida(&planos, p, rgba, ancho, alto);
Ok(imagen::respuesta(ancho as u32, alto as u32, &rgb)) Ok(imagen::respuesta(ancho as u32, alto as u32, &rgb))
} }
Ok((forma, _)) => Err(format!("El modelo contestó {forma:?} y se esperaba [1, 3, {mh}, {mw}].")), Ok((forma, _)) => Err(format!("El modelo contestó {forma:?} y se esperaba [1, 3, {}, {}].", p.ph, p.pw)),
Err(e) => { Err(e) => {
*m = None; *m = None;
crate::salud::anotar("ERROR", "ia", &format!("no se pudo rellenar: {e}")); crate::salud::anotar("ERROR", "ia", &format!("no se pudo rellenar: {e}"));
@@ -357,12 +357,12 @@ mod pruebas {
rgba[i + 3] = if hueco { 255 } else { 0 }; rgba[i + 3] = if hueco { 255 } else { 0 };
} }
} }
let (mw, mh) = relleno::tamano_modelo(w, h); let p = relleno::plan(w, h);
let t = Instant::now(); let t = Instant::now();
let (_, planos) = m.correr_forma(&[1, 4, mh, mw], relleno::entrada(&rgba, w, h, mw, mh)).unwrap(); let (_, planos) = m.correr_forma(&[1, 4, p.ph, p.pw], relleno::entrada(&rgba, w, h, p)).unwrap();
let rgb = relleno::salida(&planos, mw, mh, &rgba, w, h); let rgb = relleno::salida(&planos, p, &rgba, w, h);
let c = (h / 2 * w + w / 2) * 3; let c = (h / 2 * w + w / 2) * 3;
eprintln!("{w}×{h} (a {mw}×{mh}): {:.2} s, centro {:?}, esquina {:?}", t.elapsed().as_secs_f32(), &rgb[c..c + 3], &rgb[..3]); eprintln!("{w}×{h} ({p:?}): {:.2} s, centro {:?}, esquina {:?}", t.elapsed().as_secs_f32(), &rgb[c..c + 3], &rgb[..3]);
// fuera del hueco, lo que había // fuera del hueco, lo que había
assert_eq!(&rgb[..3], &rgba[..3]); assert_eq!(&rgb[..3], &rgba[..3]);
} }

Ver fichero

@@ -2,9 +2,14 @@
// //
// La página (lib/revelado/relleno y revelado-relleno de labstream-os) manda el trozo de la foto // La página (lib/revelado/relleno y revelado-relleno de labstream-os) manda el trozo de la foto
// alrededor de lo pintado, en 8 bits, con el hueco en el alfa (255 = rellenar). Aquí se lleva al // alrededor de lo pintado, en 8 bits, con el hueco en el alfa (255 = rellenar). Aquí se lleva al
// tamaño del modelo —el del trozo, en múltiplos de 32, con 1024 px de lado como mucho: LaMa // tamaño del modelo, se rellena, se vuelve a su tamaño y se devuelve el trozo con el hueco relleno y
// generaliza bien hasta ahí y en la CPU tarda 10 s; a 1536, 75—, se rellena, se vuelve a su tamaño y // lo demás tal cual estaba (byte a byte).
// se devuelve el trozo con el hueco relleno y lo demás tal cual estaba (byte a byte). //
// El tamaño (`plan`): LaMa hace FFT de lo que ve, y la de ONNX Runtime en la CPU solo es rápida con
// lados que son potencias de 2: 512 × 512 son 2,7 s y 704 × 448 —casi los mismos píxeles— 10 s
// (14-oct-2026). Así que el trozo se reduce lo justo para caber en 512 (o en 1024 si pasa de 768 px:
// LaMa generaliza bien hasta ahí, 10 s) y se completa hasta la potencia de 2 con su propio borde en
// espejo (sin hueco: es contexto). Lo de fuera de la foto reducida se descarta al volver.
// //
// El modelo (LaMa, Apache-2.0; la exportación ONNX de tamaño libre de g-ronimo): una entrada // El modelo (LaMa, Apache-2.0; la exportación ONNX de tamaño libre de g-ronimo): una entrada
// [1, 4, alto, ancho] —RGB 0..1 con el hueco a cero, y el hueco 0/1— y una salida [1, 3, alto, ancho] // [1, 4, alto, ancho] —RGB 0..1 con el hueco a cero, y el hueco 0/1— y una salida [1, 3, alto, ancho]
@@ -13,13 +18,23 @@
/// El lado mayor al que se rellena. /// El lado mayor al que se rellena.
pub const LADO_MAX: usize = 1024; pub const LADO_MAX: usize = 1024;
/// El tamaño al que se rellena un trozo de `ancho` × `alto`: el suyo (o reducido a LADO_MAX de lado /// Cómo entra un trozo de `ancho` × `alto` al modelo: reducido a `sw` × `sh` y completado en espejo
/// mayor), en múltiplos de 32 y con 64 como mínimo. /// hasta `pw` × `ph` (potencias de 2).
pub fn tamano_modelo(ancho: usize, alto: usize) -> (usize, usize) { #[derive(Debug, PartialEq, Eq, Clone, Copy)]
let mayor = ancho.max(alto) as f64; pub struct Plan {
let k = if mayor > LADO_MAX as f64 { LADO_MAX as f64 / mayor } else { 1.0 }; pub sw: usize,
let a32 = |v: f64| ((((v * k) / 32.0).round() as usize).max(2)) * 32; pub sh: usize,
(a32(ancho as f64), a32(alto as f64)) pub pw: usize,
pub ph: usize,
}
pub fn plan(ancho: usize, alto: usize) -> Plan {
let mayor = ancho.max(alto);
let lado = if mayor <= 768 { 512 } else { LADO_MAX };
let k = (lado as f64 / mayor as f64).min(1.0);
let sw = ((ancho as f64 * k).round() as usize).clamp(8, lado);
let sh = ((alto as f64 * k).round() as usize).clamp(8, lado);
Plan { sw, sh, pw: sw.next_power_of_two().max(64), ph: sh.next_power_of_two().max(64) }
} }
/// Remuestrea `de` (`canales` por punto, `aw` × `ah`) a `bw` × `bh` entre sus cuatro vecinos. /// Remuestrea `de` (`canales` por punto, `aw` × `ah`) a `bw` × `bh` entre sus cuatro vecinos.
@@ -46,8 +61,18 @@ pub fn remuestrear(de: &[f32], aw: usize, ah: usize, canales: usize, bw: usize,
out out
} }
/// La entrada del modelo (planos [4, mh, mw]) a partir del RGBA de la página, y el hueco a su tamaño. /// El índice en espejo de `i` dentro de 0..n (…2 1 0 1 2…n-2 n-1 n-2…).
pub fn entrada(rgba: &[u8], ancho: usize, alto: usize, mw: usize, mh: usize) -> Vec<f32> { fn espejo(i: usize, n: usize) -> usize {
if n <= 1 {
return 0;
}
let periodo = 2 * (n - 1);
let r = i % periodo;
if r < n { r } else { periodo - r }
}
/// La entrada del modelo (planos [4, ph, pw]) a partir del RGBA de la página.
pub fn entrada(rgba: &[u8], ancho: usize, alto: usize, p: Plan) -> Vec<f32> {
let n = ancho * alto; let n = ancho * alto;
let mut rgb = Vec::with_capacity(n * 3); let mut rgb = Vec::with_capacity(n * 3);
let mut hueco = Vec::with_capacity(n); let mut hueco = Vec::with_capacity(n);
@@ -55,32 +80,42 @@ pub fn entrada(rgba: &[u8], ancho: usize, alto: usize, mw: usize, mh: usize) ->
rgb.extend_from_slice(&[rgba[i * 4] as f32 / 255.0, rgba[i * 4 + 1] as f32 / 255.0, rgba[i * 4 + 2] as f32 / 255.0]); rgb.extend_from_slice(&[rgba[i * 4] as f32 / 255.0, rgba[i * 4 + 1] as f32 / 255.0, rgba[i * 4 + 2] as f32 / 255.0]);
hueco.push(if rgba[i * 4 + 3] >= 128 { 1.0 } else { 0.0 }); hueco.push(if rgba[i * 4 + 3] >= 128 { 1.0 } else { 0.0 });
} }
let rgb = remuestrear(&rgb, ancho, alto, 3, mw, mh); let rgb = remuestrear(&rgb, ancho, alto, 3, p.sw, p.sh);
let hueco = remuestrear(&hueco, ancho, alto, 1, mw, mh); let hueco = remuestrear(&hueco, ancho, alto, 1, p.sw, p.sh);
let m = mw * mh; let m = p.pw * p.ph;
let mut out = vec![0f32; 4 * m]; let mut out = vec![0f32; 4 * m];
for i in 0..m { for y in 0..p.ph {
// lo que toca el hueco, aunque sea un poco al reducir, es hueco let dentro_y = y < p.sh;
let h = if hueco[i] > 0.02 { 1.0 } else { 0.0 }; let sy = espejo(y, p.sh);
for x in 0..p.pw {
let dentro = dentro_y && x < p.sw;
let sx = espejo(x, p.sw);
let s = sy * p.sw + sx;
// lo que toca el hueco, aunque sea un poco al reducir, es hueco; el espejo, nunca
let h = if dentro && hueco[s] > 0.02 { 1.0 } else { 0.0 };
let i = y * p.pw + x;
for c in 0..3 { for c in 0..3 {
out[c * m + i] = if h > 0.0 { 0.0 } else { rgb[i * 3 + c] }; out[c * m + i] = if h > 0.0 { 0.0 } else { rgb[s * 3 + c] };
} }
out[3 * m + i] = h; out[3 * m + i] = h;
} }
}
out out
} }
/// La salida del modelo (planos [3, mh, mw], 0..1) de vuelta a `ancho` × `alto`: dentro del hueco lo /// La salida del modelo (planos [3, ph, pw], 0..1) de vuelta a `ancho` × `alto`: dentro del hueco lo
/// rellenado; fuera, el RGB que llegó, tal cual. RGB 8 bits. /// rellenado; fuera, el RGB que llegó, tal cual. RGB 8 bits.
pub fn salida(planos: &[f32], mw: usize, mh: usize, rgba: &[u8], ancho: usize, alto: usize) -> Vec<u8> { pub fn salida(planos: &[f32], p: Plan, rgba: &[u8], ancho: usize, alto: usize) -> Vec<u8> {
let m = mw * mh; let m = p.pw * p.ph;
let mut intercalado = vec![0f32; m * 3]; let mut intercalado = vec![0f32; p.sw * p.sh * 3];
for i in 0..m { for y in 0..p.sh {
for x in 0..p.sw {
for c in 0..3 { for c in 0..3 {
intercalado[i * 3 + c] = planos[c * m + i]; intercalado[(y * p.sw + x) * 3 + c] = planos[c * m + y * p.pw + x];
} }
} }
let grande = remuestrear(&intercalado, mw, mh, 3, ancho, alto); }
let grande = remuestrear(&intercalado, p.sw, p.sh, 3, ancho, alto);
let n = ancho * alto; let n = ancho * alto;
let mut out = vec![0u8; n * 3]; let mut out = vec![0u8; n * 3];
for i in 0..n { for i in 0..n {
@@ -97,26 +132,31 @@ mod pruebas {
use super::*; use super::*;
#[test] #[test]
fn el_tamano_del_modelo() { fn el_plan_reduce_lo_justo_y_completa_hasta_potencias_de_2() {
assert_eq!(tamano_modelo(500, 300), (512, 288)); assert_eq!(plan(700, 460), Plan { sw: 512, sh: 336, pw: 512, ph: 512 });
assert_eq!(tamano_modelo(2048, 1024), (1024, 512)); assert_eq!(plan(500, 300), Plan { sw: 500, sh: 300, pw: 512, ph: 512 });
assert_eq!(tamano_modelo(20, 10), (64, 64)); assert_eq!(plan(1400, 933), Plan { sw: 1024, sh: 682, pw: 1024, ph: 1024 });
let (w, h) = tamano_modelo(1600, 1067); assert_eq!(plan(200, 100), Plan { sw: 200, sh: 100, pw: 256, ph: 128 });
assert!(w <= 1024 && w % 32 == 0 && h % 32 == 0); assert_eq!(plan(20, 10), Plan { sw: 20, sh: 10, pw: 64, ph: 64 });
assert_eq!(espejo(5, 4), 1);
assert_eq!(espejo(3, 4), 3);
} }
#[test] #[test]
fn fuera_del_hueco_queda_lo_que_habia() { fn fuera_del_hueco_queda_lo_que_habia() {
// 2 × 1: el primero es hueco // 2 × 1: el primero es hueco
let rgba = [10, 20, 30, 255, 40, 50, 60, 0]; let rgba = [10, 20, 30, 255, 40, 50, 60, 0];
let e = entrada(&rgba, 2, 1, 64, 64); let p = plan(2, 1);
assert_eq!(e.len(), 4 * 64 * 64); let e = entrada(&rgba, 2, 1, p);
// a la izquierda (el hueco) el RGB va a cero y el canal del hueco a 1 let m = p.pw * p.ph;
assert_eq!(e.len(), 4 * m);
// el hueco (0,0): RGB a cero y el canal del hueco a 1; el espejo (fuera de 2 × 1) no es hueco
assert_eq!(e[0], 0.0); assert_eq!(e[0], 0.0);
assert_eq!(e[3 * 64 * 64], 1.0); assert_eq!(e[3 * m], 1.0);
assert_eq!(e[3 * m + 10], 0.0);
// lo de la red: todo gris // lo de la red: todo gris
let planos = vec![0.5f32; 3 * 64 * 64]; let planos = vec![0.5f32; 3 * m];
let s = salida(&planos, 64, 64, &rgba, 2, 1); let s = salida(&planos, p, &rgba, 2, 1);
assert_eq!(&s[..3], &[128, 128, 128]); assert_eq!(&s[..3], &[128, 128, 128]);
assert_eq!(&s[3..], &[40, 50, 60]); assert_eq!(&s[3..], &[40, 50, 60]);
} }