Labstream/labstream-desktop
1
0
Fork 0
Código Incidencias Pull Requests Paquetes Proyectos Lanzamientos Wiki Actividad

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

Explorar el Código 
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
Descargar archivo de parche Descargar archivo de diferencias Expandir todos los archivos Contraer todos los archivos

20
src-tauri/src/ia/mod.rs
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 {
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 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();
match r {
Ok((_, planos)) if planos.len() == 3 * mw * mh => {
crate::salud::anotar("INFO", "ia", &format!("relleno de {ancho}×{alto} (a {mw}×{mh}) en {:.1} s", t0.elapsed().as_secs_f32()));
let rgb = relleno::salida(&planos, mw, mh, rgba, ancho, alto);
Ok((_, planos)) if planos.len() == 3 * p.pw * p.ph => {
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, p, rgba, ancho, alto);
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) => {
*m = None;
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 };
}
}
let (mw, mh) = relleno::tamano_modelo(w, h);
let p = relleno::plan(w, h);
let t = Instant::now();
let (_, planos) = m.correr_forma(&[1, 4, mh, mw], relleno::entrada(&rgba, w, h, mw, mh)).unwrap();
let rgb = relleno::salida(&planos, mw, mh, &rgba, w, h);
let (_, planos) = m.correr_forma(&[1, 4, p.ph, p.pw], relleno::entrada(&rgba, w, h, p)).unwrap();
let rgb = relleno::salida(&planos, p, &rgba, w, h);
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
assert_eq!(&rgb[..3], &rgba[..3]);
}

122
src-tauri/src/ia/relleno.rs
Ver fichero

@@ -2,9 +2,14 @@
//
// 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
// tamaño del modelo —el del trozo, en múltiplos de 32, con 1024 px de lado como mucho: LaMa
// generaliza bien hasta ahí y en la CPU tarda 10 s; a 1536, 75—, se rellena, se vuelve a su tamaño y
// se devuelve el trozo con el hueco relleno y lo demás tal cual estaba (byte a byte).
// tamaño del modelo, se rellena, se vuelve a su tamaño y 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
// [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.
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
/// mayor), en múltiplos de 32 y con 64 como mínimo.
pub fn tamano_modelo(ancho: usize, alto: usize) -> (usize, usize) {
let mayor = ancho.max(alto) as f64;
let k = if mayor > LADO_MAX as f64 { LADO_MAX as f64 / mayor } else { 1.0 };
let a32 = |v: f64| ((((v * k) / 32.0).round() as usize).max(2)) * 32;
(a32(ancho as f64), a32(alto as f64))
/// Cómo entra un trozo de `ancho` × `alto` al modelo: reducido a `sw` × `sh` y completado en espejo
/// hasta `pw` × `ph` (potencias de 2).
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct Plan {
pub sw: usize,
pub sh: usize,
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.
@@ -46,8 +61,18 @@ pub fn remuestrear(de: &[f32], aw: usize, ah: usize, canales: usize, bw: usize,
out
}
/// La entrada del modelo (planos [4, mh, mw]) a partir del RGBA de la página, y el hueco a su tamaño.
pub fn entrada(rgba: &[u8], ancho: usize, alto: usize, mw: usize, mh: usize) -> Vec<f32> {
/// El índice en espejo de `i` dentro de 0..n (…2 1 0 1 2…n-2 n-1 n-2…).
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 mut rgb = Vec::with_capacity(n * 3);
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]);
hueco.push(if rgba[i * 4 + 3] >= 128 { 1.0 } else { 0.0 });
}
let rgb = remuestrear(&rgb, ancho, alto, 3, mw, mh);
let hueco = remuestrear(&hueco, ancho, alto, 1, mw, mh);
let m = mw * mh;
let rgb = remuestrear(&rgb, ancho, alto, 3, p.sw, p.sh);
let hueco = remuestrear(&hueco, ancho, alto, 1, p.sw, p.sh);
let m = p.pw * p.ph;
let mut out = vec![0f32; 4 * m];
for i in 0..m {
// lo que toca el hueco, aunque sea un poco al reducir, es hueco
let h = if hueco[i] > 0.02 { 1.0 } else { 0.0 };
for c in 0..3 {
out[c * m + i] = if h > 0.0 { 0.0 } else { rgb[i * 3 + c] };
for y in 0..p.ph {
let dentro_y = y < p.sh;
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 {
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
}
/// 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.
pub fn salida(planos: &[f32], mw: usize, mh: usize, rgba: &[u8], ancho: usize, alto: usize) -> Vec<u8> {
let m = mw * mh;
let mut intercalado = vec![0f32; m * 3];
for i in 0..m {
for c in 0..3 {
intercalado[i * 3 + c] = planos[c * m + i];
pub fn salida(planos: &[f32], p: Plan, rgba: &[u8], ancho: usize, alto: usize) -> Vec<u8> {
let m = p.pw * p.ph;
let mut intercalado = vec![0f32; p.sw * p.sh * 3];
for y in 0..p.sh {
for x in 0..p.sw {
for c in 0..3 {
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 mut out = vec![0u8; n * 3];
for i in 0..n {
@@ -97,26 +132,31 @@ mod pruebas {
use super::*;
#[test]
fn el_tamano_del_modelo() {
assert_eq!(tamano_modelo(500, 300), (512, 288));
assert_eq!(tamano_modelo(2048, 1024), (1024, 512));
assert_eq!(tamano_modelo(20, 10), (64, 64));
let (w, h) = tamano_modelo(1600, 1067);
assert!(w <= 1024 && w % 32 == 0 && h % 32 == 0);
fn el_plan_reduce_lo_justo_y_completa_hasta_potencias_de_2() {
assert_eq!(plan(700, 460), Plan { sw: 512, sh: 336, pw: 512, ph: 512 });
assert_eq!(plan(500, 300), Plan { sw: 500, sh: 300, pw: 512, ph: 512 });
assert_eq!(plan(1400, 933), Plan { sw: 1024, sh: 682, pw: 1024, ph: 1024 });
assert_eq!(plan(200, 100), Plan { sw: 200, sh: 100, pw: 256, ph: 128 });
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]
fn fuera_del_hueco_queda_lo_que_habia() {
// 2 × 1: el primero es hueco
let rgba = [10, 20, 30, 255, 40, 50, 60, 0];
let e = entrada(&rgba, 2, 1, 64, 64);
assert_eq!(e.len(), 4 * 64 * 64);
// a la izquierda (el hueco) el RGB va a cero y el canal del hueco a 1
let p = plan(2, 1);
let e = entrada(&rgba, 2, 1, p);
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[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
let planos = vec![0.5f32; 3 * 64 * 64];
let s = salida(&planos, 64, 64, &rgba, 2, 1);
let planos = vec![0.5f32; 3 * m];
let s = salida(&planos, p, &rgba, 2, 1);
assert_eq!(&s[..3], &[128, 128, 128]);
assert_eq!(&s[3..], &[40, 50, 60]);
}