diff --git a/CHANGELOG.md b/CHANGELOG.md index d259e8f..2de47d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,17 @@ # Changelog +## Non rilasciato — engine v2 + +Il motore impara cose nuove (la versione dell'app non cambia qui: la alza chi pubblica). `engine/versioni.json`: motore 9, quindi al primo avvio l'app aggiorna il motore da sola. Il contratto con l'interfaccia è in `docs/contratto-engine-gui.md`. + +- **FLUX.2 Klein 4B** (Apache-2.0, 4 passi) come quarto modello. Misurato su RTX 4060: Crea 5,4 s contro 6,4 s di Klein 3, zona 21 s contro 25 s, Modifica 9,6 s contro 11,5 s (`docs/bench-engine-v2.md`). Non è il predefinito. +- **LoRA multipli con pesi** su tutti i modelli (Qwen e Klein): in fila, forza da −2 a 2. Il riconoscimento dai nomi dei tensori dice per quale modello è fatto un LoRA (`src/shared/loraBase.ts`); nuovo IPC `lora.compatibili(famiglia)`. +- **Senza sfondo** (BiRefNet, MIT) e **estrai fedele** (SAM 2.1, Apache-2.0): ritagliano dai pixel della tua foto, senza ridisegnarla; i modelli si scaricano al primo uso e girano a parte. Con Klein, «sfondo trasparente» in Crea = disegno + BiRefNet. +- **Riquadri nell'Editor box** (regional prompting): Klein 3 li legge nel suo formato (JSON con le coordinate), Qwen e Klein 4B a parole. Nuovo adattatore dalla richiesta della GUI (`box`, `promptScena`, `sfondo`). +- **Catalogo con licenza e ruolo**, modelli «uncensored» marcati e scaricabili solo su richiesta: interruttore acceso, avviso accettato e licenza valida (Apache, MIT, BSD, CC-BY o Qwen Research con avviso; mai FLUX Kommunity). Un LLM uncensored (Qwen3-VL 4B abliterated) per i testi. +- Correzioni: Klein 4B ha la sua chiave VRAM nella regola «una cosa alla volta»; le parole chiave dei LoRA entrano nel prompt anche con Klein. +- Prove: `npm test` (modelli, catalogo, spazio, lora, regioni, adatta, grafi v2, grafi Qwen identici). + ## 1.1.0 — 2026-10-21 Un'app più semplice da usare: tre modelli al posto delle «qualità», una sola scheda per scrivere cosa cambiare (tutta la foto o solo una zona), le ricette dentro Modifica e il testo scritto dal Qwen3-VL 4B anche con Qwen. Il motore passa a ComfyUI 0.42.3. diff --git a/docs/bench-engine-v2.md b/docs/bench-engine-v2.md new file mode 100644 index 0000000..dfe5b98 --- /dev/null +++ b/docs/bench-engine-v2.md @@ -0,0 +1,38 @@ +# Benchmark engine v2 — prima / dopo + +Misurato il 22/10/2026 su RTX 4060 8 GB (driver 617.42), ComfyUI 0.42.3 (engine installato, motore 8), tutto 1024×1024, stesso prompt e stessa maschera di `scripts/misure_gpu.mjs`. +Grafi veri dell'app (`grafi.ts` impacchettato con esbuild). GPU libera prima di ogni giro (< 15 %), un benchmark alla volta, ComfyUI di prova sulla porta 8819 (non quello dell'app, non quello degli asset: il 8818 era fermo e vuoto). Dati completi: `docs/bench-engine-v2.json`. +Come si ripete: `node scripts/bench_serie.mjs --repo . --out --etichetta --rip 3 --porta 8819 ` con `MOTORE_OUT` = cartella di output del motore (`node scripts/misure_gpu.mjs --prova elenco` stampa le prove). + +**Prima** = Klein 3 veloce, 8 passi (il percorso rapido di prima di engine v2). **Dopo** = FLUX.2 Klein 4B FP8, 4 passi (`Flux2Scheduler` + `CFGGuider`). + +## A caldo (mediana di 3 giri dopo un giro di riscaldamento) + +| Scenario | Modello | Passi | Totale | s/passo | VRAM picco* | Peggiore | +|---|---|---:|---:|---:|---:|---:| +| Crea | Klein 3 veloce (prima) | 8 | 6,4 s | 0,471 | 6,1 GB | 6,4 s | +| Crea | **Klein 4B (dopo)** | 4 | **5,4 s** | 0,572 | 5,3 GB | 5,4 s | +| Crea | Klein 4B, 8 passi (controllo) | 8 | 8,5 s | 0,711 | 7,1 GB | 8,5 s | +| Modifica zona | Klein 3 veloce (prima) | 8 | 24,9 s | 1,873 | 7,4 GB | 25,7 s | +| Modifica zona | **Klein 4B (dopo)** | 4 | **21,3 s** | 2,258 | 6,1 GB | 21,5 s | +| Modifica tutta la foto | Klein 3 veloce (prima) | 8 | 11,5 s | 1,050 | 6,9 GB | 11,5 s | +| Modifica tutta la foto | **Klein 4B (dopo)** | 4 | **9,6 s** | 1,307 | 7,3 GB | 9,6 s | + +## A freddo (1 giro; modelli scaricati con `/free`, i file restano nella cache del disco) + +| Scenario | Modello | Passi | Totale | s/passo | VRAM picco* | +|---|---|---:|---:|---:|---:| +| Crea | Klein 3 veloce | 8 | 9,7 s | 0,468 | 6,8 GB | +| Crea | Klein 4B | 4 | 8,6 s | 0,623 | 6,9 GB | +| Crea | Klein 3 Base (CFG 3,5, riferimento) | 50 | 59,4 s | 1,006 | 6,8 GB | +| Crea | Qwen-Image 2.1 Q4_K_M (riferimento) | 40 | 162,9 s | 2,880 | 5,9 GB | + +\* picco letto con `nvidia-smi` una volta al secondo (include la cache di PyTorch): serve a vedere che nessun giro va fuori memoria, non a confrontare i megabyte fra modelli. Il primissimo avvio a freddo di Klein 4B (file mai letti dal disco) è costato 15,6 s. + +## Che cosa si tiene e che cosa no + +- **Tenuto: Klein 4B come quarto modello.** Guadagno reale ma piccolo: −16 % in Crea, −14 % in zona, −17 % in Modifica, perché ogni passo costa di più (0,57 contro 0,47 s) e a 8 passi il 4B non va più veloce di Klein 3 (8,5 contro 6,4 s). Un'immagine di Crea (teiera) e le due modifiche rosse sono pulite a occhio; **non** è stato fatto un confronto di qualità alla cieca con Klein 3. Non diventa il predefinito: il vantaggio vero è la licenza Apache-2.0, non la velocità. +- **Niente altro è stato aggiunto in nome della velocità.** Le ottimizzazioni suggerite dalla ricerca (SageAttention, `torch.compile`, cache approssimate, Nunchaku/SVDQuant) non sono state installate né misurate: richiedono pacchetti esterni nell'engine produttivo (wheel Windows, graph break con l'offload) e la ricerca stessa le classifica «esperimenti separati». Non c'è un numero a loro favore, quindi non entrano. +- Già in 1.1.0 e non rimisurato qui: il 4B che scrive i testi (185,9 → 133,5 s su «Migliora + Crea» con Qwen), vedi `docs/misure-1.1.0.md`. +- BiRefNet e SAM 2 non sono stati cronometrati: i loro nodi (`DaProdSfondo`, `DaProdSegmenta`) arrivano con il motore 9 e il motore installato qui è l'8. Girano a parte, senza il modello che disegna in VRAM. +- Un incidente di misura, scartato: un secondo `bench_serie` partito per errore in parallelo ha sporcato un primo tentativo; i dati sopra vengono da una ripetizione con un solo benchmark vivo (controllato con la lista dei processi). diff --git a/docs/contratto-engine-gui.md b/docs/contratto-engine-gui.md new file mode 100644 index 0000000..b8880ab --- /dev/null +++ b/docs/contratto-engine-gui.md @@ -0,0 +1,70 @@ +# Contratto engine ↔ GUI (engine v2) + +Chi lo usa: la GUI v2 (`feat/gui-v2`). Chi lo possiede: l'engine (`feat/engine-v2`): **`src/shared/tipi.ts`, `src/shared/modelli.ts`, `src/shared/loraBase.ts`, `src/renderer/src/api.ts`, `src/preload/index.ts` sono dell'engine**. La GUI tocca solo pagine e componenti del renderer; se serve un campo nuovo lo chiede (commento sul task), non lo aggiunge. +Tutto passa da `window.daprod` (tipo `Api` in `api.ts`); i canali IPC sono `area:azione` (`gestisci('lavori:accoda', …)` in `src/main/index.ts`). +Provato da: `npm test` (modelli, catalogo, spazio, lora, regioni, adatta, grafi_v2, confronto grafi Qwen) e `npm run typecheck`. + +## 1. Modelli (famiglie) + +`Famiglia = 'qwen' | 'klein' | 'klein-base' | 'klein4b'`. `IdModello` (`modelli.ts`) ha le stesse chiavi più `qwen-base`; `MODELLI[id]` dà nome, passi, CFG, `lora: true` per tutti. + +| Famiglia | Modello | Passi (bozza/alta/massima) | CFG | Licenza pesi | +|---|---|---|---|---| +| `qwen` | Qwen-Image 2.1 GGUF | 20/40/50 | 1 | Qwen Research (solo ricerca) | +| `klein` | FLUX 3 Klein (INT8) | 8/8/16 | 1 | FLUX Kommunity | +| `klein-base` | FLUX 3 Klein Base | 25/50/64 | 3,5 | FLUX Kommunity | +| `klein4b` | FLUX.2 Klein 4B (FP8) | 4/4/8 | 1 | **Apache-2.0** | + +`Richiesta.famiglia` (se manca vale `Impostazioni.famiglia`). Tre cose da sapere: Klein 4B usa `Flux2Scheduler` + `CFGGuider`; con Klein i testi (migliora prompt, descrivi…) li scrive il Qwen3-VL 4B; nelle opere in galleria il modello è `«FLUX.2 Klein 4B · file»` (`modelloDaOpera` lo riconosce). + +## 2. Richiesta (`lavori.accoda(q): Promise` → id dei lavori) + +Campi comuni: `modalita`, `famiglia`, `prompt`, `negativo`, `seed` (−1 = casuale), `passi`, `cfg`, `larghezza/altezza`, `megapixel`, `immagini[]` (percorsi; [0] è la foto), `quante`, `uscita`, `lora`. + +Modalità: `crea`, `modifica`, `zona` (maschera), `espandi`, `rifinisci`, `varia`, `ingrandisci`, `descrivi` (testo), `traduci`, **`sfondo`**, **`segmenta`**. + +### 2.1 LoRA multipli con pesi — `Richiesta.lora: LoraAttiva[]` +`{ file, forza }`, **l'ordine della lista è l'ordine di applicazione**, forza −2…2 (fuori scala si tronca; 0 = ignorato). Vale per **tutte** le famiglie. Il motore mette un `LoraLoaderModelOnly` per ciascuno, in catena. La parola chiave di ogni LoRA acceso entra nel prompt da sola (`conParoleLora`, ora anche con Klein). +Compatibilità: `InfoLora.famiglie: Famiglia[]` (dai nomi dei tensori, `shared/loraBase.ts`) e `InfoLora.compatibile: 'si'|'no'|'?'`. `vaBeneCon(lora, famiglia)` (shared, importabile dal renderer) dice `si|no|?` (`?` = si lascia provare). Un LoRA di Klein 3 **non** va su Klein 4B (e viceversa): sono architetture diverse. +**IPC nuovo:** `lora.compatibili(famiglia): Promise` = solo i `si` e `?` per quel modello. (`lora.elenco()` resta: tutti.) + +### 2.2 Senza sfondo +- Foto esistente: `modalita: 'sfondo'`, `immagini: [foto]` → PNG RGBA (BiRefNet, MIT). I pixel del soggetto restano quelli originali, i bordi sono ripuliti dagli aloni. Non serve il modello che disegna. +- In Crea: `trasparente: true` oppure `sfondo: 'trasparente'`. Qwen usa il suo canale alfa; **Klein (tutte) genera e poi toglie lo sfondo con BiRefNet** (la GUI può dire «Klein: sfondo tolto dopo»). +- `sfondo: { tinta: '#rrggbb' }` → una riga in più nel prompt («on a plain solid #rrggbb background»). `sfondo: 'scena'` = niente. +- Il primo uso scarica BiRefNet (~450 MB) da solo; si vede nella `fase` del lavoro («Scarico BiRefNet · 40%»). Stato: `modelli.segmentatori()`. + +### 2.3 Estrai personaggi/oggetti (SAM 2) +`modalita: 'segmenta'`, `immagini: [foto]`, `segmenta: { box?: [x0,y0,x1,y1], punti?: {x,y,positivo}[], ritaglia?: boolean }`. Coordinate **0–1 sulla foto**, origine in alto a sinistra. `ritaglia` true (predefinito) = PNG stretto attorno alla cosa (+12 px), false = tela intera col resto trasparente. Senza `box` né `punti` il lavoro fallisce con «dimmi cosa estrarre». Sempre dai pixel originali (nessuna ricreazione). Scarica SAM 2.1 small (~185 MB) al primo uso. +**IPC nuovo:** `modelli.segmentatori(): Promise` = le voci BiRefNet e SAM 2 con `presente`. + +### 2.4 Regional prompting (Editor box) — per `crea` e `modifica` +Due modi equivalenti, il motore normalizza (`adattaRichiesta`, `motore/adatta.ts`): + +1. **`box: BoxRegione[]` + `promptScena`** (quello che manda la GUI v2): `{ id, rect:[x0,y0,x1,y1] (0–1 sulla tela finale), tipo, prompt, rif?, rigidita, ordine, nascosta? }`. Le `nascosta` e quelle senza prompt non partono; l'ordine è `ordine`; `promptScena` (se c'è e ci sono riquadri validi) diventa il `prompt`. +2. `regioni: Regione[]` (`{ id, box, prompt }`) già pronte. + +Traduzione per modello: **Klein 3 / Base** → prompt «didascalia che cita ``» + righe JSON con `bbox [y_min,x_min,y_max,x_max]` su 0–1000 (formato nativo, `modifica`: `tgt_bbox`); **Qwen e Klein 4B** → la posizione a parole («at the top left of the frame»). Al massimo 16 riquadri, minimo 2 % per lato. I prompt dei riquadri vengono tradotti in inglese come quello della scena (se `traduci` è acceso). +**Limiti, detti chiari:** i riquadri sono una *guida morbida* (l'elemento può uscire un po'); `rigidita: 'blocca'` per ora vale `'decisa'`; `rif` (foto di riferimento del box) non è ancora usato dal motore. + +### 2.5 Modifica con istruzioni +`modalita: 'modifica'` (tutta la foto) e `'zona'` (solo la maschera, il resto identico al pixel): il `prompt` è l'istruzione («make the teapot red»). Famiglie Klein: le foto di riferimento sono `immagini[1..]` (``…). + +## 3. Catalogo, licenze, uncensored + +`modelli.stato(): StatoVoce[]` restituisce **tutte** le voci con: `id`, `nome`, `tipo`, `presente`, `byte`, `licenza`, `licenzaUrl`, `ruolo` (`denoiser|encoder|llm|vae|lora|segmentatore|sfondo`), `uncensored?`, `avviso?`. `modelli.scarica(id)` / `annulla(id)` / evento `su.download`. +**La guardia è nel main**, non nella UI (`motivoDiBlocco`): una voce `uncensored` si scarica solo se **tutte e tre** le condizioni valgono, e se no `modelli.scarica` va in errore con il motivo in italiano: +1. `Impostazioni.uncensored === true` (interruttore «Modelli uncensored», spento di serie); +2. `Impostazioni.uncensoredAccettato === true` (ha letto l'avviso, `AVVISO_UNCENSORED` in `catalogo.ts`); +3. **licenza valida** (`motivoLicenza`): Apache-2.0, MIT, BSD-3-Clause, CC-BY-4.0 oppure «Qwen Research» con l'avviso sul limite d'uso. FLUX Kommunity (§4(d), niente scorciatoie per togliere le protezioni), «non-commercial» non letti, licenza o link mancanti = mai. +Nessuna voce uncensored è `consigliato`: non si scarica mai da sola né in idle. Nel catalogo ci sono: encoder Qwen3-VL 8B uncensored (AtomicChat, mradermacher; Apache), LoRA `lora-uc` (Qwen Research, con avviso), LLM **Qwen3-VL 4B abliterated** (Huihui, Apache) per i testi con Klein (`Impostazioni.testoUncensored`, serve anche l'interruttore). **Non c'è un denoiser FLUX uncensored** e non verrà aggiunto. +`Impostazioni.uncensored` è un **booleano**, non l'oggetto `{abilitato, accettatoIl…}` che la GUI v2 aveva ipotizzato: la GUI deve leggere/scrivere `uncensored`, `uncensoredAccettato`, `testoUncensored` con `impostazioni.salva`. + +## 4. Coda, eventi, VRAM +Un lavoro alla volta al motore (`lavori.elenco`, `su.lavori`, `su.lavoro`). Il lavoro ha `fase` (testo in italiano, anche i download di BiRefNet/SAM2), `passo/passiTotali`, `stima`, `anteprima`, `risultati[]` (percorsi in galleria), `errore`. «Una cosa alla volta in VRAM»: cambiando famiglia (Qwen → Klein → Klein 4B) il motore scarica il modello precedente (`/free`) e la fase dice «Scarico FLUX.2 Klein 4B». `sfondo` e `segmenta` non chiedono VRAM (il nodo carica e toglie il segmentatore). + +## 5. File e cartelle da conoscere +`modelli.file()` ora include `klein4b` e `klein4bTe`. Un modello Klein 4B è «pronto» con modello + text encoder (Qwen3 4B) + VAE di FLUX.2 (`kleinVae`, lo stesso di Klein 3); voci del catalogo `klein4b-fp8`, `klein4b-te`, `klein-vae`. + +## 6. Cosa NON c'è (ancora) +`estrai.ritaglia` come IPC a parte: si usa `lavori.accoda` con `modalita: 'segmenta'`. `box.rif` e `rigidita: 'blocca'` non cambiano il risultato. Nessuna prova di qualità con GPU di `sfondo`/`segmenta` sul motore installato (i nodi `DaProdSfondo`/`DaProdSegmenta` arrivano col motore 9: al primo avvio dopo l'aggiornamento l'app aggiorna il motore da sola). diff --git a/engine/custom_nodes/DaProd-Nodi/BIREFNET_LICENSE.txt b/engine/custom_nodes/DaProd-Nodi/BIREFNET_LICENSE.txt new file mode 100644 index 0000000..485921e --- /dev/null +++ b/engine/custom_nodes/DaProd-Nodi/BIREFNET_LICENSE.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 ZhengPeng + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/engine/custom_nodes/DaProd-Nodi/__init__.py b/engine/custom_nodes/DaProd-Nodi/__init__.py index c44fb94..4e6eb45 100644 --- a/engine/custom_nodes/DaProd-Nodi/__init__.py +++ b/engine/custom_nodes/DaProd-Nodi/__init__.py @@ -10,6 +10,7 @@ import torch.nn.functional as F from .traduzione import DaProdTraduci +from .segmentazione import DaProdSegmenta, DaProdSfondo def _maschera_come(maschera, h, w): @@ -115,9 +116,13 @@ def esegui(self, image, reference, mask, strength): "DaProdRiempiZona": DaProdRiempiZona, "DaProdAccordaColori": DaProdAccordaColori, "DaProdTraduci": DaProdTraduci, + "DaProdSfondo": DaProdSfondo, + "DaProdSegmenta": DaProdSegmenta, } NODE_DISPLAY_NAME_MAPPINGS = { "DaProdRiempiZona": "DaProd · Riempi la zona", "DaProdAccordaColori": "DaProd · Accorda i colori", "DaProdTraduci": "DaProd · Traduci il prompt", + "DaProdSfondo": "DaProd · Togli lo sfondo (BiRefNet)", + "DaProdSegmenta": "DaProd · Estrai fedele (SAM 2)", } diff --git a/engine/custom_nodes/DaProd-Nodi/birefnet_compat.py b/engine/custom_nodes/DaProd-Nodi/birefnet_compat.py new file mode 100644 index 0000000..47f2740 --- /dev/null +++ b/engine/custom_nodes/DaProd-Nodi/birefnet_compat.py @@ -0,0 +1,54 @@ +# Le poche cose che il codice di BiRefNet (birefnet_modello.py, di ZhengPeng7, licenza MIT: vedi BIREFNET_LICENSE.txt) prendeva da `timm`: +# le abbiamo qui, così il motore non ha bisogno di un pacchetto in più. Sono le stesse funzioni (DropPath = profondità stocastica, spenta in inferenza). +import collections.abc +from itertools import repeat + +import torch +import torch.nn as nn +from transformers import PretrainedConfig + +trunc_normal_ = nn.init.trunc_normal_ + + +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + return tuple(repeat(x, n)) + + return parse + + +to_2tuple = _ntuple(2) + + +def drop_path(x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True): + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + +class DropPath(nn.Module): + def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): + super().__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) + + def extra_repr(self): + return f"drop_prob={round(self.drop_prob, 3):0.3f}" + + +class BiRefNetConfig(PretrainedConfig): + model_type = "SegformerForSemanticSegmentation" + + def __init__(self, bb_pretrained=False, **kwargs): + self.bb_pretrained = bb_pretrained + super().__init__(**kwargs) diff --git a/engine/custom_nodes/DaProd-Nodi/birefnet_modello.py b/engine/custom_nodes/DaProd-Nodi/birefnet_modello.py new file mode 100644 index 0000000..cc8c102 --- /dev/null +++ b/engine/custom_nodes/DaProd-Nodi/birefnet_modello.py @@ -0,0 +1,2248 @@ +### config.py + +import os +import math +from transformers import PretrainedConfig + + +class Config(PretrainedConfig): + def __init__(self) -> None: + # Compatible with the latest version of transformers + super().__init__() + + # PATH settings + self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx + + # TASK settings + self.task = ['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'][0] + self.training_set = { + 'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0], + 'COD': 'TR-COD10K+TR-CAMO', + 'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5], + 'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD', # leave DIS-VD for evaluation. + 'P3M-10k': 'TR-P3M-10k', + }[self.task] + self.prompt4loc = ['dense', 'sparse'][0] + + # Faster-Training settings + self.load_all = True + self.compile = True # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch. + # Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting. + # 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607. + # 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training. + self.precisionHigh = True + + # MODEL settings + self.ms_supervision = True + self.out_ref = self.ms_supervision and True + self.dec_ipt = True + self.dec_ipt_split = True + self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder + self.mul_scl_ipt = ['', 'add', 'cat'][2] + self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2] + self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1] + self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0] + + # TRAINING settings + self.batch_size = 4 + self.IoU_finetune_last_epochs = [ + 0, + { + 'DIS5K': -50, + 'COD': -20, + 'HRSOD': -20, + 'DIS5K+HRSOD+HRS10K': -20, + 'P3M-10k': -20, + }[self.task] + ][1] # choose 0 to skip + self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4) # DIS needs high lr to converge faster. Adapt the lr linearly + self.size = 1024 + self.num_workers = max(4, self.batch_size) # will be decrease to min(it, batch_size) at the initialization of the data_loader + + # Backbone settings + self.bb = [ + 'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 2 + 'swin_v1_t', 'swin_v1_s', # 3, 4 + 'swin_v1_b', 'swin_v1_l', # 5-bs9, 6-bs4 + 'pvt_v2_b0', 'pvt_v2_b1', # 7, 8 + 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5 + ][6] + self.lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + 'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96], + 'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64], + }[self.bb] + if self.mul_scl_ipt == 'cat': + self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection] + self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else [] + + # MODEL settings - inactive + self.lat_blk = ['BasicLatBlk'][0] + self.dec_channels_inter = ['fixed', 'adap'][0] + self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0] + self.progressive_ref = self.refine and True + self.ender = self.progressive_ref and False + self.scale = self.progressive_ref and 2 + self.auxiliary_classification = False # Only for DIS5K, where class labels are saved in `dataset.py`. + self.refine_iteration = 1 + self.freeze_bb = False + self.model = [ + 'BiRefNet', + ][0] + if self.dec_blk == 'HierarAttDecBlk': + self.batch_size = 2 ** [0, 1, 2, 3, 4][2] + + # TRAINING settings - inactive + self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4] + self.optimizer = ['Adam', 'AdamW'][1] + self.lr_decay_epochs = [1e5] # Set to negative N to decay the lr in the last N-th epoch. + self.lr_decay_rate = 0.5 + # Loss + self.lambdas_pix_last = { + # not 0 means opening this loss + # original rate -- 1 : 30 : 1.5 : 0.2, bce x 30 + 'bce': 30 * 1, # high performance + 'iou': 0.5 * 1, # 0 / 255 + 'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64) + 'mse': 150 * 0, # can smooth the saliency map + 'triplet': 3 * 0, + 'reg': 100 * 0, + 'ssim': 10 * 1, # help contours, + 'cnt': 5 * 0, # help contours + 'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4. + } + self.lambdas_cls = { + 'ce': 5.0 + } + # Adv + self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training + self.lambda_adv_d = 3. * (self.lambda_adv_g > 0) + + # PATH settings - inactive + self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis') + self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights') + self.weights = { + 'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'), + 'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]), + 'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]), + 'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]), + 'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]), + 'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]), + } + + # Callbacks - inactive + self.verbose_eval = True + self.only_S_MAE = False + self.use_fp16 = False # Bugs. It may cause nan in training. + self.SDPA_enabled = False # Bugs. Slower and errors occur in multi-GPUs + + # others + self.device = [0, 'cpu'][0] # .to(0) == .to('cuda:0') + + self.batch_size_valid = 1 + self.rand_seed = 7 + # run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f] + # with open(run_sh_file[0], 'r') as f: + # lines = f.readlines() + # self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0]) + # self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0]) + # self.val_step = [0, self.save_step][0] + + def print_task(self) -> None: + # Return task for choosing settings in shell scripts. + print(self.task) + + + +### models/backbones/pvt_v2.py + +import torch +import torch.nn as nn +from functools import partial + +from .birefnet_compat import DropPath, to_2tuple, trunc_normal_ + + +import math + +# from config import Config + +# config = Config() + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.dwconv = DWConv(hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = self.fc1(x) + x = self.dwconv(x, H, W) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class Attention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1): + super().__init__() + assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}." + + self.dim = dim + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + self.q = nn.Linear(dim, dim, bias=qkv_bias) + self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.sr_ratio = sr_ratio + if sr_ratio > 1: + self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio) + self.norm = nn.LayerNorm(dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + B, N, C = x.shape + q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + if self.sr_ratio > 1: + x_ = x.permute(0, 2, 1).reshape(B, C, H, W) + x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) + x_ = self.norm(x_) + kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + else: + kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + k, v = kv[0], kv[1] + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B, N, C) + else: + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class Block(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, + attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio) + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = x + self.drop_path(self.attn(self.norm1(x), H, W)) + x = x + self.drop_path(self.mlp(self.norm2(x), H, W)) + + return x + + +class OverlapPatchEmbed(nn.Module): + """ Image to Patch Embedding + """ + + def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + + self.img_size = img_size + self.patch_size = patch_size + self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1] + self.num_patches = self.H * self.W + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride, + padding=(patch_size[0] // 2, patch_size[1] // 2)) + self.norm = nn.LayerNorm(embed_dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x): + x = self.proj(x) + _, _, H, W = x.shape + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + + return x, H, W + + +class PyramidVisionTransformerImpr(nn.Module): + def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512], + num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0., + attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm, + depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]): + super().__init__() + self.num_classes = num_classes + self.depths = depths + + # patch_embed + self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels, + embed_dim=embed_dims[0]) + self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0], + embed_dim=embed_dims[1]) + self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1], + embed_dim=embed_dims[2]) + self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2], + embed_dim=embed_dims[3]) + + # transformer encoder + dpr = np.linspace(0, drop_path_rate, sum(depths)).tolist() # stochastic depth decay rule + cur = 0 + self.block1 = nn.ModuleList([Block( + dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[0]) + for i in range(depths[0])]) + self.norm1 = norm_layer(embed_dims[0]) + + cur += depths[0] + self.block2 = nn.ModuleList([Block( + dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[1]) + for i in range(depths[1])]) + self.norm2 = norm_layer(embed_dims[1]) + + cur += depths[1] + self.block3 = nn.ModuleList([Block( + dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[2]) + for i in range(depths[2])]) + self.norm3 = norm_layer(embed_dims[2]) + + cur += depths[2] + self.block4 = nn.ModuleList([Block( + dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[3]) + for i in range(depths[3])]) + self.norm4 = norm_layer(embed_dims[3]) + + # classification head + # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity() + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def init_weights(self, pretrained=None): + if isinstance(pretrained, str): + logger = 1 + #load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger) + + def reset_drop_path(self, drop_path_rate): + dpr = np.linspace(0, drop_path_rate, sum(self.depths)).tolist() + cur = 0 + for i in range(self.depths[0]): + self.block1[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[0] + for i in range(self.depths[1]): + self.block2[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[1] + for i in range(self.depths[2]): + self.block3[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[2] + for i in range(self.depths[3]): + self.block4[i].drop_path.drop_prob = dpr[cur + i] + + def freeze_patch_emb(self): + self.patch_embed1.requires_grad = False + + @torch.jit.ignore + def no_weight_decay(self): + return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better + + def get_classifier(self): + return self.head + + def reset_classifier(self, num_classes, global_pool=''): + self.num_classes = num_classes + self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + + def forward_features(self, x): + B = x.shape[0] + outs = [] + + # stage 1 + x, H, W = self.patch_embed1(x) + for i, blk in enumerate(self.block1): + x = blk(x, H, W) + x = self.norm1(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 2 + x, H, W = self.patch_embed2(x) + for i, blk in enumerate(self.block2): + x = blk(x, H, W) + x = self.norm2(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 3 + x, H, W = self.patch_embed3(x) + for i, blk in enumerate(self.block3): + x = blk(x, H, W) + x = self.norm3(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 4 + x, H, W = self.patch_embed4(x) + for i, blk in enumerate(self.block4): + x = blk(x, H, W) + x = self.norm4(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + return outs + + # return x.mean(dim=1) + + def forward(self, x): + x = self.forward_features(x) + # x = self.head(x) + + return x + + +class DWConv(nn.Module): + def __init__(self, dim=768): + super(DWConv, self).__init__() + self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) + + def forward(self, x, H, W): + B, N, C = x.shape + x = x.transpose(1, 2).view(B, C, H, W).contiguous() + x = self.dwconv(x) + x = x.flatten(2).transpose(1, 2) + + return x + + +def _conv_filter(state_dict, patch_size=16): + """ convert patch embedding weight from manual patchify + linear proj to conv""" + out_dict = {} + for k, v in state_dict.items(): + if 'patch_embed.proj.weight' in k: + v = v.reshape((v.shape[0], 3, patch_size, patch_size)) + out_dict[k] = v + + return out_dict + + +class pvt_v2_b0(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b0, self).__init__( + patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +class pvt_v2_b1(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b1, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +class pvt_v2_b2(PyramidVisionTransformerImpr): + def __init__(self, in_channels=3, **kwargs): + super(pvt_v2_b2, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels) + +class pvt_v2_b3(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b3, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +class pvt_v2_b4(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b4, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + +class pvt_v2_b5(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b5, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +### models/backbones/swin_v1.py + +# -------------------------------------------------------- +# Swin Transformer +# Copyright (c) 2021 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Written by Ze Liu, Yutong Lin, Yixuan Wei +# -------------------------------------------------------- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +import numpy as np +from .birefnet_compat import DropPath, to_2tuple, trunc_normal_ + +# from config import Config + + +# config = Config() + + +class Mlp(nn.Module): + """ Multilayer perceptron.""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """ Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ Forward function. + + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B_, N, C) + else: + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww, Wh*Ww, nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ Swin Transformer Block. + + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """ Patch Merging Layer + + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop=0., + attn_drop=0., + drop_path=0., + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + # Turn int to torch.tensor for the compatiability with torch.compile in PyTorch 2.5. + Hp = torch.ceil(torch.tensor(H) / self.window_size).to(torch.int64) * self.window_size + Wp = torch.ceil(torch.tensor(W) / self.window_size).to(torch.int64) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)).to(x.dtype) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + + Args: + patch_size (int): Patch token size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_channels = in_channels + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """ Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + pretrain_img_size=224, + patch_size=4, + in_channels=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop_rate=0., + attn_drop_rate=0., + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + use_checkpoint=False): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] + + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = np.linspace(0, drop_path_rate, sum(depths)).tolist() # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2 ** i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint) + self.layers.append(layer) + + num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f'norm{i_layer}' + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + + def forward(self, x): + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic') + x = (x + absolute_pos_embed) # B Wh*Ww C + + outs = []#x.contiguous()] + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + if i in self.out_indices: + norm_layer = getattr(self, f'norm{i}') + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + + return tuple(outs) + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + +def swin_v1_t(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_s(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_b(): + model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12) + return model + +def swin_v1_l(): + model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12) + return model + + + +### models/modules/deform_conv.py + +import torch +import torch.nn as nn +from torchvision.ops import deform_conv2d + + +class DeformableConv2d(nn.Module): + def __init__(self, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + bias=False): + + super(DeformableConv2d, self).__init__() + + assert type(kernel_size) == tuple or type(kernel_size) == int + + kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size) + self.stride = stride if type(stride) == tuple else (stride, stride) + self.padding = padding + + self.offset_conv = nn.Conv2d(in_channels, + 2 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.offset_conv.weight, 0.) + nn.init.constant_(self.offset_conv.bias, 0.) + + self.modulator_conv = nn.Conv2d(in_channels, + 1 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.modulator_conv.weight, 0.) + nn.init.constant_(self.modulator_conv.bias, 0.) + + self.regular_conv = nn.Conv2d(in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=bias) + + def forward(self, x): + #h, w = x.shape[2:] + #max_offset = max(h, w)/4. + + offset = self.offset_conv(x)#.clamp(-max_offset, max_offset) + modulator = 2. * torch.sigmoid(self.modulator_conv(x)) + + x = deform_conv2d( + input=x, + offset=offset, + weight=self.regular_conv.weight, + bias=self.regular_conv.bias, + padding=self.padding, + mask=modulator, + stride=self.stride, + ) + return x + + + + +### utils.py + +import torch.nn as nn + + +def build_act_layer(act_layer): + if act_layer == 'ReLU': + return nn.ReLU(inplace=True) + elif act_layer == 'SiLU': + return nn.SiLU(inplace=True) + elif act_layer == 'GELU': + return nn.GELU() + + raise NotImplementedError(f'build_act_layer does not support {act_layer}') + + +def build_norm_layer(dim, + norm_layer, + in_format='channels_last', + out_format='channels_last', + eps=1e-6): + layers = [] + if norm_layer == 'BN': + if in_format == 'channels_last': + layers.append(to_channels_first()) + layers.append(nn.BatchNorm2d(dim)) + if out_format == 'channels_last': + layers.append(to_channels_last()) + elif norm_layer == 'LN': + if in_format == 'channels_first': + layers.append(to_channels_last()) + layers.append(nn.LayerNorm(dim, eps=eps)) + if out_format == 'channels_first': + layers.append(to_channels_first()) + else: + raise NotImplementedError( + f'build_norm_layer does not support {norm_layer}') + return nn.Sequential(*layers) + + +class to_channels_first(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 3, 1, 2) + + +class to_channels_last(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 2, 3, 1) + + + +### dataset.py + +_class_labels_TR_sorted = ( + 'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, ' + 'BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, ' + 'CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, ' + 'Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, ' + 'Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, ' + 'Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, ' + 'KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, ' + 'Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, ' + 'OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, ' + 'RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, ' + 'ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, ' + 'Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, ' + 'TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, ' + 'UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht' +) +class_labels_TR_sorted = _class_labels_TR_sorted.split(', ') + + +### models/backbones/build_backbones.py + +import torch +import torch.nn as nn +from collections import OrderedDict +from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights +# from models.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5 +# from models.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l +# from config import Config + + +config = Config() + +def build_backbone(bb_name, pretrained=True, params_settings=''): + if bb_name == 'vgg16': + bb_net = list(vgg16(pretrained=VGG16_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:4], 'conv2': bb_net[4:9], 'conv3': bb_net[9:16], 'conv4': bb_net[16:23]})) + elif bb_name == 'vgg16bn': + bb_net = list(vgg16_bn(pretrained=VGG16_BN_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:6], 'conv2': bb_net[6:13], 'conv3': bb_net[13:23], 'conv4': bb_net[23:33]})) + elif bb_name == 'resnet50': + bb_net = list(resnet50(pretrained=ResNet50_Weights.DEFAULT if pretrained else None).children()) + bb = nn.Sequential(OrderedDict({'conv1': nn.Sequential(*bb_net[0:3]), 'conv2': bb_net[4], 'conv3': bb_net[5], 'conv4': bb_net[6]})) + else: + bb = eval('{}({})'.format(bb_name, params_settings)) + if pretrained: + bb = load_weights(bb, bb_name) + return bb + +def load_weights(model, model_name): + save_model = torch.load(config.weights[model_name], map_location='cpu') + model_dict = model.state_dict() + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()} + # to ignore the weights with mismatched size when I modify the backbone itself. + if not state_dict: + save_model_keys = list(save_model.keys()) + sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()} + if not state_dict or not sub_item: + print('Weights are not successully loaded. Check the state dict of weights file.') + return None + else: + print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item)) + model_dict.update(state_dict) + model.load_state_dict(model_dict) + return model + + + +### models/modules/decoder_blocks.py + +import torch +import torch.nn as nn +# from models.aspp import ASPP, ASPPDeformable +# from config import Config + + +# config = Config() + + +class BasicDecBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicDecBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.relu_in = nn.ReLU(inplace=True) + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + def forward(self, x): + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + + +class ResBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=None, inter_channels=64): + super(ResBlk, self).__init__() + if out_channels is None: + out_channels = in_channels + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.relu_in = nn.ReLU(inplace=True) + + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + _x = self.conv_resi(x) + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + _x + + + +### models/modules/lateral_blocks.py + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial + +# from config import Config + + +# config = Config() + + +class BasicLatBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicLatBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + x = self.conv(x) + return x + + + +### models/modules/aspp.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +# from models.deform_conv import DeformableConv2d +# from config import Config + + +# config = Config() + + +class _ASPPModule(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding, dilation): + super(_ASPPModule, self).__init__() + self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, dilation=dilation, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPP(nn.Module): + def __init__(self, in_channels=64, out_channels=None, output_stride=16): + super(ASPP, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + if output_stride == 16: + dilations = [1, 6, 12, 18] + elif output_stride == 8: + dilations = [1, 12, 24, 36] + else: + raise NotImplementedError + + self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0]) + self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1]) + self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2]) + self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x2 = self.aspp2(x) + x3 = self.aspp3(x) + x4 = self.aspp4(x) + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, x2, x3, x4, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + +##################### Deformable +class _ASPPModuleDeformable(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding): + super(_ASPPModuleDeformable, self).__init__() + self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPPDeformable(nn.Module): + def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]): + super(ASPPDeformable, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + + self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0) + self.aspp_deforms = nn.ModuleList([ + _ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes + ]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x_aspp_deforms = [aspp_deform(x) for aspp_deform in self.aspp_deforms] + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, *x_aspp_deforms, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + + +### models/refinement/refiner.py + +import torch +import torch.nn as nn +from collections import OrderedDict +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.models import vgg16, vgg16_bn +from torchvision.models import resnet50 + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.ing import * +# from models.stem_layer import StemLayer + + +class RefinerPVTInChannels4(nn.Module): + def __init__(self, in_channels=3+1): + super(RefinerPVTInChannels4, self).__init__() + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, params_settings='in_channels=4') + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Refiner(nn.Module): + def __init__(self, in_channels=3+1): + super(Refiner, self).__init__() + self.config = Config() + self.epoch = 1 + self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + self.bb = build_backbone(self.config.bb) + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + x = self.stem_layer(x) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval('BasicDecBlk') + LateralBlock = eval('BasicLatBlk') + + self.decoder_block4 = DecoderBlock(channels[0], channels[1]) + self.decoder_block3 = DecoderBlock(channels[1], channels[2]) + self.decoder_block2 = DecoderBlock(channels[2], channels[3]) + self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0)) + + def forward(self, features): + x, x1, x2, x3, x4 = features + outs = [] + p4 = self.decoder_block4(x4) + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + p3 = self.decoder_block3(_p3) + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + p2 = self.decoder_block2(_p2) + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision: + outs.append(self.conv_ms_spvn_4(p4)) + outs.append(self.conv_ms_spvn_3(p3)) + outs.append(self.conv_ms_spvn_2(p2)) + outs.append(p1_out) + return outs + + +class RefUNet(nn.Module): + # Refinement + def __init__(self, in_channels=3+1): + super(RefUNet, self).__init__() + self.encoder_1 = nn.Sequential( + nn.Conv2d(in_channels, 64, 3, 1, 1), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_2 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_3 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_4 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True) + ##### + self.decoder_5 = nn.Sequential( + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + ##### + self.decoder_4 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_3 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_2 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_1 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1) + + self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) + + def forward(self, x): + outs = [] + if isinstance(x, list): + x = torch.cat(x, dim=1) + hx = x + + hx1 = self.encoder_1(hx) + hx2 = self.encoder_2(hx1) + hx3 = self.encoder_3(hx2) + hx4 = self.encoder_4(hx3) + + hx = self.decoder_5(self.pool4(hx4)) + hx = torch.cat((self.upscore2(hx), hx4), 1) + + d4 = self.decoder_4(hx) + hx = torch.cat((self.upscore2(d4), hx3), 1) + + d3 = self.decoder_3(hx) + hx = torch.cat((self.upscore2(d3), hx2), 1) + + d2 = self.decoder_2(hx) + hx = torch.cat((self.upscore2(d2), hx1), 1) + + d1 = self.decoder_1(hx) + + x = self.conv_d0(d1) + outs.append(x) + return outs + + + +### models/stem_layer.py + +import torch.nn as nn +# from utils import build_act_layer, build_norm_layer + + +class StemLayer(nn.Module): + r""" Stem layer of InternImage + Args: + in_channels (int): number of input channels + out_channels (int): number of output channels + act_layer (str): activation layer + norm_layer (str): normalization layer + """ + + def __init__(self, + in_channels=3+1, + inter_channels=48, + out_channels=96, + act_layer='GELU', + norm_layer='BN'): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, + inter_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm1 = build_norm_layer( + inter_channels, norm_layer, 'channels_first', 'channels_first' + ) + self.act = build_act_layer(act_layer) + self.conv2 = nn.Conv2d(inter_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm2 = build_norm_layer( + out_channels, norm_layer, 'channels_first', 'channels_first' + ) + + def forward(self, x): + x = self.conv1(x) + x = self.norm1(x) + x = self.act(x) + x = self.conv2(x) + x = self.norm2(x) + return x + + +### models/birefnet.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +from kornia.filters import laplacian +from transformers import PreTrainedModel +from einops import rearrange + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.aspp import ASPP, ASPPDeformable +# from models.ing import * +# from models.refiner import Refiner, RefinerPVTInChannels4, RefUNet +# from models.stem_layer import StemLayer +from .birefnet_compat import BiRefNetConfig + + +def image2patches(image, grid_h=2, grid_w=2, patch_ref=None, transformation='b c (hg h) (wg w) -> (b hg wg) c h w'): + if patch_ref is not None: + grid_h, grid_w = image.shape[-2] // patch_ref.shape[-2], image.shape[-1] // patch_ref.shape[-1] + patches = rearrange(image, transformation, hg=grid_h, wg=grid_w) + return patches + +def patches2image(patches, grid_h=2, grid_w=2, patch_ref=None, transformation='(b hg wg) c h w -> b c (hg h) (wg w)'): + if patch_ref is not None: + grid_h, grid_w = patch_ref.shape[-2] // patches[0].shape[-2], patch_ref.shape[-1] // patches[0].shape[-1] + image = rearrange(patches, transformation, hg=grid_h, wg=grid_w) + return image + +class BiRefNet( + PreTrainedModel +): + config_class = BiRefNetConfig + def __init__(self, bb_pretrained=True, config=BiRefNetConfig()): + super(BiRefNet, self).__init__(config) + bb_pretrained = config.bb_pretrained + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained) + + channels = self.config.lateral_channels_in_collection + + if self.config.auxiliary_classification: + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.cls_head = nn.Sequential( + nn.Linear(channels[0], len(class_labels_TR_sorted)) + ) + + if self.config.squeeze_block: + self.squeeze_module = nn.Sequential(*[ + eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0]) + for _ in range(eval(self.config.squeeze_block.split('_x')[1])) + ]) + + self.decoder = Decoder(channels) + + if self.config.ender: + self.dec_end = nn.Sequential( + nn.Conv2d(1, 16, 3, 1, 1), + nn.Conv2d(16, 1, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + # refine patch-level segmentation + if self.config.refine: + if self.config.refine == 'itself': + self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + else: + self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1')) + + if self.config.freeze_bb: + # Freeze the backbone... + print(self.named_parameters()) + for key, value in self.named_parameters(): + if 'bb.' in key and 'refiner.' not in key: + value.requires_grad = False + + self.post_init() + + def forward_enc(self, x): + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + if self.config.mul_scl_ipt == 'cat': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1) + elif self.config.mul_scl_ipt == 'add': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = x1 + F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True) + x2 = x2 + F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True) + x3 = x3 + F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True) + x4 = x4 + F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True) + class_preds = self.cls_head(self.avgpool(x4).view(x4.shape[0], -1)) if self.training and self.config.auxiliary_classification else None + if self.config.cxt: + x4 = torch.cat( + ( + *[ + F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True), + ][-len(self.config.cxt):], + x4 + ), + dim=1 + ) + return (x1, x2, x3, x4), class_preds + + def forward_ori(self, x): + ########## Encoder ########## + (x1, x2, x3, x4), class_preds = self.forward_enc(x) + if self.config.squeeze_block: + x4 = self.squeeze_module(x4) + ########## Decoder ########## + features = [x, x1, x2, x3, x4] + if self.training and self.config.out_ref: + features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5)) + scaled_preds = self.decoder(features) + return scaled_preds, class_preds + + def forward(self, x): + scaled_preds, class_preds = self.forward_ori(x) + class_preds_lst = [class_preds] + return [scaled_preds, class_preds_lst] if self.training else scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval(self.config.dec_blk) + LateralBlock = eval(self.config.lat_blk) + + if self.config.dec_ipt: + self.split = self.config.dec_ipt_split + N_dec_ipt = 64 + DBlock = SimpleConvs + ic = 64 + ipt_cha_opt = 1 + self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic) + else: + self.split = None + + self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1]) + self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2]) + self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]) + self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0)) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + + if self.config.out_ref: + _N = 16 + self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + + self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + def forward(self, features): + if self.training and self.config.out_ref: + outs_gdt_pred = [] + outs_gdt_label = [] + x, x1, x2, x3, x4, gdt_gt = features + else: + x, x1, x2, x3, x4 = features + outs = [] + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=x4, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1) + p4 = self.decoder_block4(x4) + m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p4_gdt = self.gdt_convs_4(p4) + if self.training: + # >> GT: + m4_dia = m4 + gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_4) + # >> Pred: + gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt) + outs_gdt_pred.append(gdt_pred_4) + gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid() + # >> Finally: + p4 = p4 * gdt_attn_4 + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p3, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1) + p3 = self.decoder_block3(_p3) + m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p3_gdt = self.gdt_convs_3(p3) + if self.training: + # >> GT: + # m3 --dilation--> m3_dia + # G_3^gt * m3_dia --> G_3^m, which is the label of gradient + m3_dia = m3 + gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_3) + # >> Pred: + # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx + # F_3^G --sigmoid--> A_3^G + gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt) + outs_gdt_pred.append(gdt_pred_3) + gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid() + # >> Finally: + # p3 = p3 * A_3^G + p3 = p3 * gdt_attn_3 + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p2, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1) + p2 = self.decoder_block2(_p2) + m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p2_gdt = self.gdt_convs_2(p2) + if self.training: + # >> GT: + m2_dia = m2 + gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_2) + # >> Pred: + gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt) + outs_gdt_pred.append(gdt_pred_2) + gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid() + # >> Finally: + p2 = p2 * gdt_attn_2 + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1) + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision and self.training: + outs.append(m4) + outs.append(m3) + outs.append(m2) + outs.append(p1_out) + return outs if not (self.config.out_ref and self.training) else ([outs_gdt_pred, outs_gdt_label], outs) + + +class SimpleConvs(nn.Module): + def __init__( + self, in_channels: int, out_channels: int, inter_channels=64 + ) -> None: + super().__init__() + self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1) + + def forward(self, x): + return self.conv_out(self.conv1(x)) diff --git a/engine/custom_nodes/DaProd-Nodi/segmentazione.py b/engine/custom_nodes/DaProd-Nodi/segmentazione.py new file mode 100644 index 0000000..8e03799 --- /dev/null +++ b/engine/custom_nodes/DaProd-Nodi/segmentazione.py @@ -0,0 +1,252 @@ +# Togli sfondo e estrai fedele, per DaProd Image. +# - DaProdSfondo: BiRefNet (ZhengPeng7, MIT) separa il soggetto dallo sfondo; l'uscita è la foto ORIGINALE con il canale alfa (i pixel del soggetto non si ridisegnano). +# - DaProdSegmenta: SAM 2.1 (Meta, Apache-2.0, via transformers) ritaglia la cosa indicata da un riquadro o da punti, sempre dai pixel originali. +# Tutte e due: modello da models/segment/ (si scarica dall'app al primo uso), caricato sulla GPU solo per il tempo del lavoro +# e poi tolto, così il modello che disegna ha di nuovo tutta la VRAM. +import json +import os + +import torch +import torch.nn.functional as F + + +def _dispositivo(): + try: + import comfy.model_management as mm + + return mm.get_torch_device() + except Exception: + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def _libera(dispositivo): + if getattr(dispositivo, "type", "cpu") == "cuda": + torch.cuda.empty_cache() + + +# ---------- pezzi puri (si provano senza modelli) ---------- + + +def pulisci_colori(rgb, alfa, solido=0.9): + """rgb [B,H,W,3], alfa [B,H,W] (0..1). Ai bordi i colori dei pixel semi-trasparenti hanno dentro un po' di sfondo (aloni chiari o scuri): + si rifanno con i colori del soggetto vicino (push-pull, come DaProdRiempiZona). I pixel solidi non si toccano.""" + from . import riempi + + da_rifare = (alfa < solido).float() + if float(da_rifare.sum()) == 0.0 or float(da_rifare.sum()) == float(da_rifare.numel()): + return rgb + rifatto = riempi(rgb, da_rifare) + m = (alfa >= solido).float().unsqueeze(-1) + return rgb * m + rifatto * (1.0 - m) + + +def componi_rgba(rgb, alfa, pulisci=True): + """foto originale + alfa -> immagine [B,H,W,4]""" + rgb = rgb[..., :3] + if pulisci: + rgb = pulisci_colori(rgb, alfa) + return torch.cat([rgb, alfa.unsqueeze(-1)], dim=-1).clamp(0, 1) + + +def ritaglia_su_alfa(rgba, soglia=0.02, margine=12): + """ritaglia [B,H,W,4] attorno ai pixel con alfa > soglia (+ margine); se non c'è niente, restituisce com'è. Ritorna (immagine, (x0,y0,x1,y1))""" + a = rgba[0, :, :, 3] + ys, xs = torch.where(a > soglia) + h, w = a.shape + if ys.numel() == 0: + return rgba, (0, 0, w, h) + y0 = max(0, int(ys.min()) - margine) + y1 = min(h, int(ys.max()) + 1 + margine) + x0 = max(0, int(xs.min()) - margine) + x1 = min(w, int(xs.max()) + 1 + margine) + return rgba[:, y0:y1, x0:x1, :], (x0, y0, x1, y1) + + +def leggi_box(testo, larghezza, altezza): + """'[x0,y0,x1,y1]' in 0..1 -> [x0,y0,x1,y1] in pixel, dentro la foto; None se vuoto o non valido""" + if not testo or not testo.strip(): + return None + try: + b = json.loads(testo) + x0, y0, x1, y1 = [float(v) for v in b] + except Exception: + return None + x0, x1 = sorted((x0, x1)) + y0, y1 = sorted((y0, y1)) + x0, x1 = max(0.0, x0) * larghezza, min(1.0, x1) * larghezza + y0, y1 = max(0.0, y0) * altezza, min(1.0, y1) * altezza + if x1 - x0 < 2 or y1 - y0 < 2: + return None + return [x0, y0, x1, y1] + + +def leggi_punti(testo, larghezza, altezza): + """'[{"x":0.4,"y":0.5,"positivo":true}, …]' (0..1) -> ([[x,y],…] in pixel, [1|0,…])""" + if not testo or not testo.strip(): + return [], [] + try: + lista = json.loads(testo) + except Exception: + return [], [] + punti, etichette = [], [] + for p in lista if isinstance(lista, list) else []: + try: + x = min(1.0, max(0.0, float(p["x"]))) * larghezza + y = min(1.0, max(0.0, float(p["y"]))) * altezza + except Exception: + continue + punti.append([x, y]) + etichette.append(0 if p.get("positivo") is False else 1) + return punti, etichette + + +# ---------- BiRefNet ---------- + +_MEDIA = (0.485, 0.456, 0.406) +_DEV = (0.229, 0.224, 0.225) + + +def _carica_birefnet(cartella, dispositivo): + from safetensors.torch import load_file + + from .birefnet_compat import BiRefNetConfig + from .birefnet_modello import BiRefNet + + modello = BiRefNet(bb_pretrained=False, config=BiRefNetConfig()) + stato = load_file(os.path.join(cartella, "model.safetensors")) + # alcuni salvataggi hanno il prefisso del modulo + stato = {k[7:] if k.startswith("module.") else k: v for k, v in stato.items()} + mancano, inutili = modello.load_state_dict(stato, strict=False) + mancano = [k for k in mancano if "num_batches_tracked" not in k] + if mancano: + raise RuntimeError("BiRefNet: nel file mancano %d pesi (primo: %s)" % (len(mancano), mancano[0])) + return modello.eval().to(dispositivo).half() if dispositivo.type == "cuda" else modello.eval().to(dispositivo) + + +def alfa_birefnet(immagine, cartella, lato=1024): + """immagine [B,H,W,3] 0..1 -> alfa [B,H,W] 0..1""" + dispositivo = _dispositivo() + modello = _carica_birefnet(cartella, dispositivo) + try: + b, h, w, _ = immagine.shape + x = immagine[..., :3].permute(0, 3, 1, 2).float() + x = F.interpolate(x, size=(lato, lato), mode="bilinear", align_corners=False) + media = torch.tensor(_MEDIA).view(1, 3, 1, 1) + dev = torch.tensor(_DEV).view(1, 3, 1, 1) + x = ((x - media) / dev).to(dispositivo) + if dispositivo.type == "cuda": + x = x.half() + fuori = [] + with torch.inference_mode(): + for i in range(b): + p = modello(x[i : i + 1])[-1].sigmoid().float().cpu() + fuori.append(F.interpolate(p, size=(h, w), mode="bilinear", align_corners=False)[:, 0]) + return torch.cat(fuori, dim=0).clamp(0, 1) + finally: + del modello + _libera(dispositivo) + + +# ---------- SAM 2 ---------- + + +def alfa_sam2(immagine, cartella, box, punti, etichette): + """una foto [1,H,W,3] -> maschera [1,H,W] 0..1 della cosa indicata (box in pixel e/o punti in pixel)""" + import numpy as np + from PIL import Image + from transformers import Sam2Model, Sam2Processor + + dispositivo = _dispositivo() + elaboratore = Sam2Processor.from_pretrained(cartella) + modello = Sam2Model.from_pretrained(cartella).eval().to(dispositivo) + try: + foto = Image.fromarray((immagine[0, :, :, :3].clamp(0, 1) * 255).round().byte().cpu().numpy()) + arg = {} + if box is not None: + arg["input_boxes"] = [[list(box)]] + if punti: + arg["input_points"] = [[[list(p) for p in punti]]] + arg["input_labels"] = [[list(etichette)]] + dati = elaboratore(images=foto, return_tensors="pt", **arg).to(dispositivo) + with torch.inference_mode(): + out = modello(**dati, multimask_output=False) + maschere = elaboratore.post_process_masks(out.pred_masks.cpu(), dati["original_sizes"].cpu(), binarize=False)[0] + m = maschere.float() + # [oggetti, scelte, H, W] -> la prima scelta del primo oggetto; i valori sono logit: si porta a 0..1 con la sigmoide + while m.dim() > 2: + m = m[0] + return torch.sigmoid(m).unsqueeze(0).clamp(0, 1) + finally: + del modello + _libera(dispositivo) + + +# ---------- i nodi ---------- + + +class DaProdSfondo: + CATEGORY = "DaProd" + FUNCTION = "esegui" + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "mask") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "modello": ("STRING", {"default": "birefnet"}), + "pulisci_bordi": ("BOOLEAN", {"default": True}), + "ritaglia": ("BOOLEAN", {"default": False}), + } + } + + def esegui(self, image, modello, pulisci_bordi, ritaglia): + import folder_paths + + cartella = os.path.join(folder_paths.models_dir, "segment", modello) + if not os.path.isfile(os.path.join(cartella, "model.safetensors")): + raise RuntimeError("manca BiRefNet in " + cartella) + alfa = alfa_birefnet(image, cartella) + rgba = componi_rgba(image, alfa, pulisci_bordi) + if ritaglia: + rgba, _ = ritaglia_su_alfa(rgba) + return (rgba, rgba[..., 3]) + + +class DaProdSegmenta: + CATEGORY = "DaProd" + FUNCTION = "esegui" + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "mask") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "modello": ("STRING", {"default": "sam2.1-hiera-small"}), + "box": ("STRING", {"default": ""}), + "punti": ("STRING", {"default": ""}), + "ritaglia": ("BOOLEAN", {"default": True}), + "margine": ("INT", {"default": 12, "min": 0, "max": 512}), + "pulisci_bordi": ("BOOLEAN", {"default": True}), + } + } + + def esegui(self, image, modello, box, punti, ritaglia, margine, pulisci_bordi): + import folder_paths + + cartella = os.path.join(folder_paths.models_dir, "segment", modello) + if not os.path.isfile(os.path.join(cartella, "model.safetensors")): + raise RuntimeError("manca SAM 2 in " + cartella) + _, h, w, _ = image.shape + b = leggi_box(box, w, h) + p, e = leggi_punti(punti, w, h) + if b is None and not p: + raise RuntimeError("dimmi cosa estrarre: un riquadro o dei punti") + alfa = alfa_sam2(image[:1], cartella, b, p, e) + rgba = componi_rgba(image[:1], alfa, pulisci_bordi) + if ritaglia: + rgba, _ = ritaglia_su_alfa(rgba, margine=margine) + return (rgba, rgba[..., 3]) diff --git a/engine/versioni.json b/engine/versioni.json index 2370f2e..6e4bb4b 100644 --- a/engine/versioni.json +++ b/engine/versioni.json @@ -1,5 +1,5 @@ { - "motore": 8, + "motore": 9, "python": "3.13", "comfyui": { "commit": "edb3db00fd88f059e748c8a400391b5eab76266b", diff --git a/package.json b/package.json index efe4f87..b4dd86c 100644 --- a/package.json +++ b/package.json @@ -14,6 +14,7 @@ "build": "electron-vite build", "start": "electron-vite preview", "typecheck": "tsc --noEmit -p tsconfig.node.json && tsc --noEmit -p tsconfig.web.json", + "test": "node test/modelli.mjs && node test/catalogo.mjs && node test/spazio.mjs && node test/lora.mjs && node test/regioni.mjs && node test/adatta.mjs && node test/grafi_v2.mjs && node test/confronta_grafi.mjs", "prendi-uv": "node scripts/prendi-uv.mjs", "dist": "npm run prendi-uv && npm run build && electron-builder --win --x64", "prova": "node test/prove.mjs", diff --git a/scripts/bench_serie.mjs b/scripts/bench_serie.mjs new file mode 100644 index 0000000..efccb91 --- /dev/null +++ b/scripts/bench_serie.mjs @@ -0,0 +1,138 @@ +#!/usr/bin/env node +// Benchmark riproducibile di DaProdImage: ripete le prove di scripts/misure_gpu.mjs con semi diversi (così ComfyUI non risponde dalla cache), +// scarta un giro di riscaldamento per le prove «a caldo», e riassume con mediana e peggiore. Una sola variabile per volta: si lancia con lo stesso +// motore (porta 8819, avviato a parte) prima e dopo ogni modifica e si confrontano i due file. +// +// Uso: +// node scripts/bench_serie.mjs --repo --out --etichetta base [--rip 3] [--seed0 1000] [--porta 8819] prova1 prova2 … +// node scripts/bench_serie.mjs --prova elenco (stampa i nomi delle prove di misure_gpu.mjs) +// +// Le prove che finiscono in «-freddo» (e «libera», i testi «testo-…», «migliora-crea-…») girano una volta sola, senza riscaldamento. +// Prima di ogni giro si aspetta che la GPU sia ferma (< 15 %): se un altro programma la usa, il giro lo dice (`gpu_occupata`) e non finisce nella mediana. +// Scrive /serie-.json e stampa una tabella in Markdown. +import { spawnSync } from 'node:child_process' +import { mkdirSync, writeFileSync } from 'node:fs' +import { dirname, join, resolve } from 'node:path' +import { fileURLToPath } from 'node:url' + +const QUI = dirname(fileURLToPath(import.meta.url)) +const args = process.argv.slice(2) +const opt = (n, d) => { + const i = args.indexOf('--' + n) + if (i < 0) return d + const v = args[i + 1] + return v && !v.startsWith('--') ? v : true +} +const REPO = resolve(String(opt('repo', join(QUI, '..')))) +const OUT = resolve(String(opt('out', join(QUI, '..', 'bench-out')))) +const ETICHETTA = String(opt('etichetta', 'prova')) +const RIP = Number(opt('rip', 3)) +const SEED0 = Number(opt('seed0', 1000)) +const PORTA = String(opt('porta', 8819)) +const flag = new Set(['repo', 'out', 'etichetta', 'rip', 'seed0', 'porta', 'prova']) +const prove = [] +for (let i = 0; i < args.length; i++) { + if (args[i].startsWith('--')) { + if (flag.has(args[i].slice(2))) i++ + continue + } + prove.push(args[i]) +} +if (opt('prova', '') === 'elenco') { + const r = spawnSync('node', [join(REPO, 'scripts', 'misure_gpu.mjs'), '--prova', 'elenco', '--porta', PORTA, '--out', OUT], { encoding: 'utf8' }) + console.log(r.stdout) + process.exit(0) +} +if (!prove.length) { + console.error('nessuna prova: vedi --prova elenco') + process.exit(2) +} +mkdirSync(OUT, { recursive: true }) + +const dorme = (ms) => new Promise((r) => setTimeout(r, ms)) +function smi() { + const r = spawnSync('nvidia-smi', ['--query-gpu=memory.used,utilization.gpu,clocks.sm,temperature.gpu', '--format=csv,noheader,nounits'], { encoding: 'utf8', windowsHide: true }) + const p = (r.stdout || '').trim().split(',').map((x) => Number(x.trim())) + return { vram_mb: p[0], gpu_pct: p[1], clock_mhz: p[2], temp_c: p[3] } +} +async function aspettaGpuFerma() { + for (let i = 0; i < 30; i++) { + if (smi().gpu_pct < 15) return true + await dorme(2000) + } + return false +} +function mediana(v) { + const s = [...v].sort((a, b) => a - b) + const m = Math.floor(s.length / 2) + return s.length % 2 ? s[m] : (s[m - 1] + s[m]) / 2 +} +const arr = (x, d = 2) => (x === undefined || x === null || Number.isNaN(x) ? null : Math.round(x * 10 ** d) / 10 ** d) + +function giro(prova, seed) { + const r = spawnSync('node', [join(REPO, 'scripts', 'misure_gpu.mjs'), '--porta', PORTA, '--out', OUT, '--prova', prova, '--seed', String(seed)], { + encoding: 'utf8', env: { ...process.env, MOTORE_OUT: process.env.MOTORE_OUT || '' }, maxBuffer: 64 * 1024 * 1024 + }) + let j = null + try { + j = JSON.parse(r.stdout) + } catch { + /* niente */ + } + return j || { errore: (r.stderr || r.stdout || 'nessun output').slice(0, 300) } +} + +const risultati = [] +const SOLO_UNA = /-freddo$|^libera$|^testo-|^migliora-crea-/ +for (const prova of prove) { + const giri = [] + const una = SOLO_UNA.test(prova) + if (!una) { + await aspettaGpuFerma() + giro(prova, SEED0 - 1) // riscaldamento: carica i modelli, non conta + } + const n = una ? 1 : RIP + for (let i = 0; i < n; i++) { + const libera = await aspettaGpuFerma() + const prima = smi() + const r = giro(prova, SEED0 + i) + const rr = r.crea || r // migliora-crea: il giro finale + giri.push({ + seed: SEED0 + i, + totale_s: arr(r.totale_s, 1), + campionatore_s: arr(r.sampler?.secondi, 2), + s_per_passo: arr(r.sampler?.s_per_passo, 3), + passi: r.sampler?.passi, + picco_vram_mb: r.picco_vram_mb, + gpu_media_pct: r.gpu_media_pct, + gpu_occupata: !libera, + prima, + errore: r.errore + }) + void rr + } + const buoni = giri.filter((g) => !g.errore && !g.gpu_occupata && g.totale_s !== null) + const ris = { + prova, + giri, + mediana: buoni.length ? { totale_s: arr(mediana(buoni.map((g) => g.totale_s)), 1), s_per_passo: arr(mediana(buoni.filter((g) => g.s_per_passo !== null).map((g) => g.s_per_passo)), 3), picco_vram_mb: Math.max(...buoni.map((g) => g.picco_vram_mb || 0)) } : null, + peggiore_s: buoni.length ? Math.max(...buoni.map((g) => g.totale_s)) : null, + giri_validi: buoni.length, + giri_scartati: giri.length - buoni.length + } + risultati.push(ris) + console.error(`# ${prova}: ${buoni.length}/${giri.length} giri validi, mediana ${ris.mediana?.totale_s ?? '-'} s`) +} + +const macchina = { gpu: spawnSync('nvidia-smi', ['--query-gpu=name,memory.total,driver_version', '--format=csv,noheader'], { encoding: 'utf8' }).stdout.trim(), quando: new Date().toISOString(), repo: REPO, rip: RIP, seed0: SEED0 } +const gitHead = spawnSync('git', ['-C', REPO, 'rev-parse', '--short', 'HEAD'], { encoding: 'utf8' }).stdout.trim() +macchina.commit = gitHead +writeFileSync(join(OUT, `serie-${ETICHETTA}.json`), JSON.stringify({ etichetta: ETICHETTA, macchina, risultati }, null, 2)) + +console.log(`\n### ${ETICHETTA} (${gitHead}, ${macchina.gpu}, ${RIP} giri a caldo, mediana)\n`) +console.log('| Prova | Totale | s/passo | Passi | VRAM picco | Peggiore | Giri validi |') +console.log('|---|---:|---:|---:|---:|---:|---:|') +for (const r of risultati) { + const g = r.giri.find((x) => x.passi) || {} + console.log(`| ${r.prova} | ${r.mediana?.totale_s ?? '-'} s | ${r.mediana?.s_per_passo ?? '-'} | ${g.passi ?? '-'} | ${r.mediana ? (r.mediana.picco_vram_mb / 1024).toFixed(1) + ' GB' : '-'} | ${r.peggiore_s ?? '-'} s | ${r.giri_validi}/${r.giri.length} |`) +} diff --git a/scripts/chiavi_safetensors.mjs b/scripts/chiavi_safetensors.mjs new file mode 100644 index 0000000..8c22230 --- /dev/null +++ b/scripts/chiavi_safetensors.mjs @@ -0,0 +1,21 @@ +// Stampa le chiavi (raggruppate) dell'intestazione di un .safetensors. Uso: node scripts/chiavi_safetensors.mjs [filtro-regex] +import { openSync, readSync, closeSync } from 'node:fs' +const [, , file, filtro] = process.argv +const fd = openSync(file, 'r') +const b8 = Buffer.alloc(8) +readSync(fd, b8, 0, 8, 0) +const n = Number(b8.readBigUInt64LE(0)) +const b = Buffer.alloc(n) +readSync(fd, b, 0, n, 8) +closeSync(fd) +const h = JSON.parse(b.toString('utf8')) +const chiavi = Object.keys(h).filter((k) => k !== '__metadata__') +const re = filtro ? new RegExp(filtro) : null +const gruppi = new Map() +for (const k of chiavi) { + if (re && !re.test(k)) continue + const g = k.replace(/\.\d+\./g, '.N.').replace(/^(\w+)\.\d+/, '$1.N') + gruppi.set(g, (gruppi.get(g) || 0) + 1) +} +console.log('tensori:', chiavi.length, 'meta:', JSON.stringify(h.__metadata__ || {}).slice(0, 300)) +for (const [g, c] of [...gruppi].slice(0, 80)) console.log(String(c).padStart(5), g, JSON.stringify(h[chiavi.find((k) => k.replace(/\.\d+\./g, '.N.').replace(/^(\w+)\.\d+/, '$1.N') === g)]?.shape)) diff --git a/scripts/chiavi_safetensors_url.mjs b/scripts/chiavi_safetensors_url.mjs new file mode 100644 index 0000000..481939d --- /dev/null +++ b/scripts/chiavi_safetensors_url.mjs @@ -0,0 +1,15 @@ +// Come chiavi_safetensors.mjs ma da un URL (Range): legge solo l'intestazione. Uso: node scripts/chiavi_safetensors_url.mjs +const url = process.argv[2] +const r8 = await fetch(url, { headers: { Range: 'bytes=0-7' }, redirect: 'follow' }) +const n = Number(Buffer.from(await r8.arrayBuffer()).readBigUInt64LE(0)) +const r = await fetch(url, { headers: { Range: `bytes=8-${8 + n - 1}` }, redirect: 'follow' }) +const h = JSON.parse(Buffer.from(await r.arrayBuffer()).toString('utf8')) +const chiavi = Object.keys(h).filter((k) => k !== '__metadata__') +const gruppi = new Map() +for (const k of chiavi) { + const g = k.replace(/\.\d+\./g, '.N.').replace(/^(\w+)\.\d+/, '$1.N') + if (!gruppi.has(g)) gruppi.set(g, { c: 0, forma: JSON.stringify(h[k].shape), dtype: h[k].dtype }) + gruppi.get(g).c++ +} +console.log('tensori:', chiavi.length, JSON.stringify(h.__metadata__ || {}).slice(0, 200)) +for (const [g, v] of [...gruppi].slice(0, 60)) console.log(String(v.c).padStart(5), g, v.forma, v.dtype) diff --git a/scripts/leggi_template.mjs b/scripts/leggi_template.mjs new file mode 100644 index 0000000..9b494b4 --- /dev/null +++ b/scripts/leggi_template.mjs @@ -0,0 +1,15 @@ +// Stampa il grafo dei nodi di un template ComfyUI (workflow JSON dell'interfaccia): tipo, valori dei widget, collegamenti. Uso: node scripts/leggi_template.mjs +import { readFileSync } from 'node:fs' +const j = JSON.parse(readFileSync(process.argv[2], 'utf8')) +const nodi = (j.nodes || j.definitions?.subgraphs?.[0]?.nodes || []) +const sub = j.definitions?.subgraphs || [] +const stampa = (lista, pref = '') => { + for (const n of lista) { + console.log(`${pref}${n.id} ${n.type}${n.title ? ' «' + n.title + '»' : ''} ${JSON.stringify(n.widgets_values ?? '')}`) + } +} +stampa(j.nodes || []) +for (const s of sub) { + console.log('--- sottografo', s.name) + stampa(s.nodes || [], ' ') +} diff --git a/scripts/misure_gpu.mjs b/scripts/misure_gpu.mjs index 536112b..511f2dd 100644 --- a/scripts/misure_gpu.mjs +++ b/scripts/misure_gpu.mjs @@ -34,17 +34,17 @@ const SEME = Number(opt('seed', 424242)) const ETICHETTA = String(opt('etichetta', '')) mkdirSync(OUT, { recursive: true }) -// ---- grafi.ts com'è ---- -function caricaGrafi() { - let s = readFileSync(join(REPO, 'src', 'main', 'motore', 'grafi.ts'), 'utf8') - s = s.replace(/^import type .*? from '@shared\/tipi'\r?\n/gm, '') - s = s.replace(/^import \{ modelloTurbo \} from '@shared\/tipi'\r?\n/m, 'const modelloTurbo = (file) => /turbo/i.test(file)\n') - s = s.replace("constructor(readonly memoria: Memoria = 'normale') {}", "memoria: Memoria = 'normale'\n constructor(memoria: Memoria = 'normale') { this.memoria = memoria }") - const p = join(tmpdir(), 'misure-gpu-grafi.ts') - writeFileSync(p, s) - return import(pathToFileURL(p).href) +// ---- grafi.ts impacchettato con esbuild (importa regioni.ts e @shared: lo strip dei tipi di Node non basta) ---- +async function caricaGrafi() { + const { buildSync } = await import('esbuild') + const fuori = join(tmpdir(), 'misure-gpu-grafi.mjs') + buildSync({ entryPoints: [join(REPO, 'src', 'main', 'motore', 'grafi.ts')], bundle: true, format: 'esm', platform: 'node', outfile: fuori, alias: { '@shared': join(REPO, 'src', 'shared') }, logLevel: 'error' }) + return import(pathToFileURL(fuori).href + '?v=' + Date.now()) } -const { costruisci } = await caricaGrafi() +const { costruisci: costruisciGrezzo } = await caricaGrafi() +// FLUX.2 Klein 4B ha i suoi file: modello FP8 e Qwen3 4B di solo testo (il VAE è quello di FLUX.2, lo stesso di Klein 3) +const FILE_4B = { kleinModello: 'flux-2-klein-4b-fp8.safetensors', kleinEncoder: 'qwen_3_4b.safetensors' } +const costruisci = (q, i, ...resto) => costruisciGrezzo(q, q.famiglia === 'klein4b' ? { ...i, ...FILE_4B } : i, ...resto) const { SISTEMA_MIGLIORA } = await (async () => { // la stringa di SISTEMA_MIGLIORA sta nel renderer (azioni.ts, ha import di React): la si legge dal sorgente const t = readFileSync(join(REPO, 'src', 'renderer', 'src', 'azioni.ts'), 'utf8') @@ -261,6 +261,12 @@ prova('klein-mod-16', async () => esegui(grafoModifica('klein', 16, false), 'Kle // riferimenti: catena vs multi prova('klein-rif-catena-8', async () => esegui(grafoModifica('klein', 8, true), 'Klein 8 passi · Modifica con riferimento · ReferenceLatent in catena')) prova('klein-rif-multi-8', async () => esegui(conMultiLatent(grafoModifica('klein', 8, true)), 'Klein 8 passi · Modifica con riferimento · ReferenceMultiLatent')) +// FLUX.2 Klein 4B (distillato, 4 passi, Flux2Scheduler + CFGGuider) +prova('klein4b-crea-4-freddo', async () => { await libera(); return esegui(grafoCrea('klein4b', 4), 'Klein 4B 4 passi · Crea · a freddo') }) +prova('klein4b-crea-4', async () => esegui(grafoCrea('klein4b', 4), 'Klein 4B 4 passi · Crea · caldo')) +prova('klein4b-crea-8', async () => esegui(grafoCrea('klein4b', 8), 'Klein 4B 8 passi · Crea · caldo')) +prova('klein4b-zona-4', async () => esegui(grafoZona('klein4b', 4), 'Klein 4B 4 passi · zona · caldo')) +prova('klein4b-mod-4', async () => esegui(grafoModifica('klein4b', 4, false), 'Klein 4B 4 passi · Modifica tutta la foto')) // Klein Base prova('base-crea-50-freddo', async () => { await libera(); return esegui(grafoCrea('klein-base', 50), 'Klein Base 50 passi · Crea · a freddo') }) prova('base-crea-50', async () => esegui(grafoCrea('klein-base', 50), 'Klein Base 50 passi · Crea · caldo')) diff --git a/scripts/salva_alberi_hf.mjs b/scripts/salva_alberi_hf.mjs new file mode 100644 index 0000000..06ea442 --- /dev/null +++ b/scripts/salva_alberi_hf.mjs @@ -0,0 +1,25 @@ +// Salva gli alberi dei repo Hugging Face usati dal catalogo (solo percorso e dimensione) in test/campioni-hf/, per provare `scegli` senza rete. +// Uso: node scripts/salva_alberi_hf.mjs (rete) +import { mkdirSync, writeFileSync, readFileSync } from 'node:fs' +import { dirname, join } from 'node:path' +import { fileURLToPath, pathToFileURL } from 'node:url' + +const QUI = dirname(fileURLToPath(import.meta.url)) +const { CATALOGO } = await import(pathToFileURL(join(QUI, '..', 'src', 'main', 'installa', 'catalogo.ts')).href) +const dest = join(QUI, '..', 'test', 'campioni-hf') +mkdirSync(dest, { recursive: true }) +const visti = new Set() +for (const v of CATALOGO) { + if (!v.repo) continue + const chiave = `${v.repo}@${v.revisione || 'main'}` + if (visti.has(chiave)) continue + visti.add(chiave) + const r = await fetch(`https://huggingface.co/api/models/${v.repo}/tree/${v.revisione || 'main'}?recursive=true`) + if (!r.ok) { + console.log('NO', chiave, r.status) + continue + } + const albero = (await r.json()).filter((f) => f.type === 'file').map((f) => ({ type: 'file', path: f.path, size: f.size, lfs: f.lfs ? { size: f.lfs.size } : undefined })) + writeFileSync(join(dest, v.repo.replace('/', '__') + '.json'), JSON.stringify({ repo: v.repo, revisione: v.revisione || 'main', file: albero }, null, 1)) + console.log('ok', chiave, albero.length, 'file') +} diff --git a/src/main/idle.ts b/src/main/idle.ts index 8481a9a..98ebfa4 100644 --- a/src/main/idle.ts +++ b/src/main/idle.ts @@ -70,7 +70,7 @@ const TOLTI = ['qwen-image-2.1-UC-Q5_K_M.gguf', 'qwen3vl_8b_bf16.safetensors'] export function fileInutilizzati(): FileInutile[] { const imp = impostazioni() // i file di Klein (scelti a mano o scaricati dall'app) sono in uso quanto quelli di Qwen - const klein = ['klein-int8', 'klein-bf16', 'klein-base-int8', 'klein-base-bf16', 'klein-te', 'klein-vae'].map((id) => risoltoSalvato(id)?.file || '') + const klein = ['klein-int8', 'klein-bf16', 'klein-base-int8', 'klein-base-bf16', 'klein-te', 'klein-vae', 'klein4b-fp8', 'klein4b-te', 'llm-huihui-4b'].map((id) => risoltoSalvato(id)?.file || '') const inUso = new Set([imp.modello, imp.encoder, imp.vae, imp.upscaler, imp.loraTurbo, imp.kleinModello, imp.kleinBase, imp.kleinEncoder, imp.kleinVae, ...klein].filter(Boolean).map((x) => x.toLowerCase())) const noti = new Set(statoModelli().filter((v) => !v.link).map((v) => (risoltoSalvato(v.id)?.file || v.file).toLowerCase()).filter(Boolean)) const mmproj = (n: string): boolean => /mmproj|\.mmproj\./i.test(n) diff --git a/src/main/impostazioni.ts b/src/main/impostazioni.ts index 259882e..6111a34 100644 --- a/src/main/impostazioni.ts +++ b/src/main/impostazioni.ts @@ -44,7 +44,10 @@ const PREDEFINITE: Impostazioni = { unaCosaAllaVolta: true, potenziaKlein: true, testoCon4B: true, // misurato in T12 (docs/misure-1.1.0.md): Migliora prompt + Crea con Qwen passa da 186 s a 134 s (encoder 8B int8) o da 446 s a 150 s (8B GGUF) - scaricaDopo: 120 + scaricaDopo: 120, + uncensored: false, + uncensoredAccettato: false, + testoUncensored: false } let attuali: Impostazioni | null = null @@ -95,6 +98,8 @@ export function impostazioni(): Impostazioni { export function salvaImpostazioni(modifiche: Partial): Impostazioni { const nuove = { ...impostazioni(), ...modifiche, ...allineaScelta(modifiche) } + // spento l'interruttore, l'LLM uncensored non scrive più i testi (resta sul disco, se l'hai scaricato) + if (modifiche.uncensored === false) nuove.testoUncensored = false if (modifiche.predefiniti) nuove.predefiniti = { ...impostazioni().predefiniti, ...modifiche.predefiniti } attuali = nuove scrivi(nuove) diff --git a/src/main/index.ts b/src/main/index.ts index 84876ee..2ff6601 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -5,7 +5,8 @@ import { appendFileSync, copyFileSync, existsSync, readFileSync, readdirSync, rm import { basename, extname, join, normalize } from 'node:path' import { pathToFileURL } from 'node:url' import { randomUUID } from 'node:crypto' -import type { Download, FiltroGalleria, Impostazioni, InfoLora, Richiesta } from '@shared/tipi' +import type { Download, Famiglia, FiltroGalleria, Impostazioni, InfoLora, Richiesta } from '@shared/tipi' +import { vaBeneCon } from '@shared/loraBase' import { DATI, ENTRATA, MINIATURE, RADICE, TEMP, assicura } from './percorsi' import { impostazioni, salvaImpostazioni } from './impostazioni' import { annullaSetup, cercaModelliEsistenti, controlloSistema, installa, motoreInstallato, scaricaVoce, statoModelli, type OpzioniSetup } from './installa/installatore' @@ -195,6 +196,8 @@ function registraIpc(): void { // modelli gestisci('modelli:stato', () => statoModelli()) + // BiRefNet (togli sfondo) e SAM 2 (estrai fedele): presenti o da scaricare (si scaricano da soli al primo uso) + gestisci('modelli:segmentatori', () => statoModelli().filter((v) => v.ruolo === 'sfondo' || v.ruolo === 'segmentatore')) gestisci('modelli:file', () => { const c = impostazioni().cartellaModelli const leggi = (sotto: string, ext: RegExp): string[] => { @@ -218,7 +221,9 @@ function registraIpc(): void { upscaler: leggi('upscale_models', /\.(pth|safetensors)$/i), klein: tutti.filter(isKlein), kleinTe: enc.filter(isKleinTe), - kleinVae: vae.filter(isKleinVae) + kleinVae: vae.filter(isKleinVae), + klein4b: tutti.filter((n) => /klein.*4b|4b.*klein/i.test(n)), + klein4bTe: enc.filter((n) => /^qwen_?3_4b/i.test(n)) } }) gestisci('modelli:inutilizzati', () => fileInutilizzati()) @@ -317,6 +322,8 @@ function registraIpc(): void { // LoRA gestisci('lora:elenco', () => elencoLora()) + // i LoRA che si possono usare con questo modello: compatibili ('si') e da provare ('?'), non quelli di un altro modello + gestisci('lora:compatibili', (famiglia: Famiglia) => elencoLora().filter((l) => vaBeneCon(l, famiglia) !== 'no')) gestisci('lora:aggiorna', (file: string, m: Partial) => aggiornaLora(file, m)) gestisci('lora:importa', (percorsi: string[]) => importaLora(percorsi)) gestisci('lora:elimina', (file: string) => eliminaLora(file)) diff --git a/src/main/installa/catalogo.ts b/src/main/installa/catalogo.ts index 40fe9b0..5f021d1 100644 --- a/src/main/installa/catalogo.ts +++ b/src/main/installa/catalogo.ts @@ -6,8 +6,24 @@ const TURBO = 'AtomicChat/Qwen-Image-2.1-Turbo-GGUF' const ENC_AC = 'AtomicChat/Qwen-Image-2.1-Turbo-Uncensored-GGUF' const ENC_MR = 'mradermacher/Qwen3-VL-8B-Instruct-Uncensored-V2-GGUF' const KLEIN = 'black-forest-labs/flux-3-image-klein-comfy' +// revisioni fissate (controllate il 22/10/2026): il download è riproducibile anche se il repo cambia dopo +const KLEIN_REV = 'feef60b0aaf63a4773ace96636ebbcf43b881312' +const KLEIN4B_FP8 = 'black-forest-labs/FLUX.2-klein-4b-fp8' +const KLEIN4B_FP8_REV = '5b4408e59397a4a37ccb46afe426d8ed86379441' +const KLEIN4B_COMFY = 'Comfy-Org/vae-text-encorder-for-flux-klein-4b' +const KLEIN4B_COMFY_REV = '87cfba3d4b98662fb80f23921bb37bf887c97c58' +const BIREFNET = 'ZhengPeng7/BiRefNet' +const BIREFNET_REV = 'e2bf8e4460fc8fa32bba5ea4d94b3233d367b0e4' +const SAM2 = 'facebook/sam2.1-hiera-small' +const SAM2_REV = 'ee5bba1d82bb8749febdf90f45e84b687142ba03' +const HUIHUI4B = 'ahmed22xa/Huihui-Qwen3-VL-4B-Instruct-abliterated-comfy' +const HUIHUI4B_REV = '6d6fc98f9bfa783dfa4f143804525742cb5dad62' -export const CATALOGO: VoceCatalogo[] = [ +const LIC_QWEN = 'https://huggingface.co/Qwen/Qwen-Image-2.1/raw/main/LICENSE' +const LIC_KOMMUNITY = 'https://huggingface.co/black-forest-labs/flux-3-image-klein/blob/main/LICENSE.md' +const LIC_APACHE = 'https://www.apache.org/licenses/LICENSE-2.0' + +const GREZZO: VoceCatalogo[] = [ // --- modello di diffusione (uno solo serve; solo Q4, Q6 e Q8) --- { id: 'q4km', tipo: 'diffusione', quant: 'Q4', nome: 'Qwen-Image-2.1 Uncensored · Q4', @@ -78,27 +94,62 @@ export const CATALOGO: VoceCatalogo[] = [ accanto: ['^config\\.json$', '^generation_config\\.json$', '^tokenizer_config\\.json$', '^source\\.spm$', '^target\\.spm$', '^vocab\\.json$'] }, // --- FLUX 3 Image Klein (Black Forest Labs): il repo «comfy» ha bf16, int8 («convrot») e NVFP4 per Klein (veloce, 8 passi) e Klein Base (50 passi, CFG). - // I nomi dei file vengono dai workflow ufficiali di ComfyUI; il nome esatto lo chiede installa/risolvi.ts a Hugging Face. Per una scheda da 6-8 GB va bene l'INT8. - // L'encoder è un Qwen3-VL 4B (diverso dal nostro 8B); il VAE è quello di FLUX.2. --- + // I percorsi esatti vengono dall'albero del repo (controllato il 22/10/2026, revisione KLEIN_REV): la prima espressione di ogni lista è il nome esatto, + // le altre sono ripieghi per i nomi vecchi (flux3_image_klein_sd_…) e per eventuali rinomine; l'ultima è larga ma esclude VAE, encoder, upsampler e altri formati. + // Per una scheda da 6-8 GB va bene l'INT8. L'encoder è un Qwen3-VL 4B (diverso dal nostro 8B); il VAE è quello di FLUX.2. --- ...([ - ['klein-int8', 'FLUX 3 Klein · veloce · INT8', 'Il più veloce: 8 passi senza CFG. L\'INT8 sta nelle schede da 8 GB.', 5_700_000_000, ['^flux3_image_klein_sd_int8', '^(?=.*klein)(?!.*base)(?!.*(vae|qwen|text_enc|bf16|nvfp4|fp8))(?=.*int8)']], - ['klein-bf16', 'FLUX 3 Klein · veloce · BF16', 'Stessa qualità dell\'originale, ma pesa il doppio: serve molta RAM.', 11_300_000_000, ['^flux3_image_klein_sd(_bf16)?\\.safetensors$', '^(?=.*klein)(?!.*base)(?!.*(vae|qwen|text_enc|int8|nvfp4|fp8))(?=.*bf16)']], - ['klein-base-int8', 'FLUX 3 Klein Base · INT8', 'La versione «normale»: 50 passi con CFG 3,5, la qualità piena (3 volte più lenta di Klein).', 5_700_000_000, ['^flux3_image_klein_base_int8', '^(?=.*klein)(?=.*base)(?!.*(vae|qwen|text_enc|bf16|nvfp4|fp8))(?=.*int8)']], - ['klein-base-bf16', 'FLUX 3 Klein Base · BF16', 'Klein Base non quantizzato.', 11_300_000_000, ['^flux3_image_klein_base(_bf16)?\\.safetensors$', '^(?=.*klein)(?=.*base)(?!.*(vae|qwen|text_enc|int8|nvfp4|fp8))(?=.*bf16)']] + ['klein-int8', 'FLUX 3 Klein · veloce · INT8', 'Il più veloce: 8 passi senza CFG. L\'INT8 sta nelle schede da 8 GB.', 5_800_000_000, ['^flux3-image-klein-int8-convrot\\.safetensors$', '^flux3_image_klein_sd_int8', '^(?=.*klein)(?!.*base)(?!.*(vae|qwen|text_enc|upsampl|bf16|nvfp4|fp8))(?=.*int8)']], + ['klein-bf16', 'FLUX 3 Klein · veloce · BF16', 'Stessa qualità dell\'originale, ma pesa il doppio: serve molta RAM.', 11_100_000_000, ['^flux3-image-klein-bf16\\.safetensors$', '^flux3_image_klein_sd(_bf16)?\\.safetensors$', '^(?=.*klein)(?!.*base)(?!.*(vae|qwen|text_enc|upsampl|int8|nvfp4|fp8))(?=.*bf16)']], + ['klein-base-int8', 'FLUX 3 Klein Base · INT8', 'La versione «normale»: 50 passi con CFG 3,5, la qualità piena (3 volte più lenta di Klein).', 5_800_000_000, ['^flux3-image-klein-base-int8-convrot\\.safetensors$', '^flux3_image_klein_base_int8', '^(?=.*klein)(?=.*base)(?!.*(vae|qwen|text_enc|upsampl|bf16|nvfp4|fp8))(?=.*int8)']], + ['klein-base-bf16', 'FLUX 3 Klein Base · BF16', 'Klein Base non quantizzato.', 11_100_000_000, ['^flux3-image-klein-base-bf16\\.safetensors$', '^flux3_image_klein_base(_bf16)?\\.safetensors$', '^(?=.*klein)(?=.*base)(?!.*(vae|qwen|text_enc|upsampl|int8|nvfp4|fp8))(?=.*bf16)']] ] as const).map(([id, nome, descrizione, byte, cerca]): VoceCatalogo => ({ - id, tipo: 'klein', nome, descrizione, file: '', cartella: 'diffusion_models', url: '', byte, circa: true, repo: KLEIN, cerca: [...cerca] + id, tipo: 'klein', nome, descrizione, file: '', cartella: 'diffusion_models', url: '', byte, circa: true, repo: KLEIN, cerca: [...cerca], revisione: KLEIN_REV })), { id: 'klein-te', tipo: 'klein-te', nome: 'Text encoder di FLUX 3 Klein (Qwen3-VL 4B)', descrizione: 'Legge il prompt per Klein. È un Qwen3-VL da 4 miliardi di parametri, diverso dal nostro encoder 8B (che resta per migliora prompt, Idee, Estrai…).', - file: '', cartella: 'text_encoders', url: '', byte: 4_500_000_000, circa: true, - repo: KLEIN, cerca: ['^qwen3vl_4b_int8', '^(?=.*(qwen3|qwen_3))(?=.*vl)(?=.*4b)(?!.*(prompt|upsampl|9b))', '^(?=.*(qwen3|qwen_3))(?=.*4b)(?!.*(prompt|upsampl|9b))'] + file: '', cartella: 'text_encoders', url: '', byte: 4_900_000_000, circa: true, + repo: KLEIN, revisione: KLEIN_REV, cerca: ['^qwen3vl_4b_int8_convrot\\.safetensors$', '^qwen3vl_4b_int8', '^(?=.*(qwen3|qwen_3))(?=.*vl)(?=.*4b)(?!.*(prompt|upsampl|9b))', '^(?=.*(qwen3|qwen_3))(?=.*4b)(?!.*(prompt|upsampl|9b))'] }, { id: 'klein-vae', tipo: 'klein-vae', nome: 'VAE di FLUX 3 Klein (FLUX.2)', descrizione: 'Il decoder delle immagini di Klein (piccolo).', - file: '', cartella: 'vae', url: '', byte: 170_000_000, circa: true, - repo: KLEIN, cerca: ['^flux2-vae\\.safetensors$', '^(?=.*flux)(?=.*vae)', '^(?=.*vae)(?!.*qwen)'] + file: '', cartella: 'vae', url: '', byte: 336_211_292, circa: true, + repo: KLEIN, revisione: KLEIN_REV, cerca: ['^flux2-vae\\.safetensors$', '^(?=.*flux)(?=.*vae)', '^(?=.*vae)(?!.*qwen)'] + }, + // --- FLUX.2 Klein 4B (Black Forest Labs, gennaio 2026): licenza Apache-2.0, 4 passi senza CFG. Il denoiser è l'FP8 ufficiale (4 GB); + // l'encoder è un Qwen3 4B di solo testo (quello dei workflow di ComfyUI); il VAE è lo stesso di Klein 3 (flux2-vae, voce `klein-vae`). --- + { + id: 'klein4b-fp8', tipo: 'klein', nome: 'FLUX.2 Klein 4B · FP8', + descrizione: 'Klein da 4 miliardi di parametri: 4 passi senza CFG, licenza Apache-2.0 (uso commerciale libero). Non fa le modifiche con tante foto bene come Klein 3.', + file: '', cartella: 'diffusion_models', url: '', byte: 4_070_624_520, circa: true, + repo: KLEIN4B_FP8, revisione: KLEIN4B_FP8_REV, cerca: ['^flux-2-klein-4b-fp8\\.safetensors$', '^(?=.*klein)(?=.*4b)(?!.*base)(?=.*fp8)'] + }, + { + id: 'klein4b-te', tipo: 'klein-te', nome: 'Text encoder di FLUX.2 Klein 4B (Qwen3 4B)', + descrizione: 'Il Qwen3 4B di solo testo che legge il prompt per FLUX.2 Klein 4B. Pesa 8 GB e lavora dalla RAM.', + file: '', cartella: 'text_encoders', url: '', byte: 8_044_982_048, circa: true, + repo: KLEIN4B_COMFY, revisione: KLEIN4B_COMFY_REV, cerca: ['^qwen_3_4b\\.safetensors$'] + }, + // --- togli sfondo e estrai fedele: si scaricano al primo uso, girano in una fase a parte (pochi centinaia di MB di VRAM) e poi escono dalla scheda --- + { + id: 'sfondo-birefnet', tipo: 'sfondo', nome: 'BiRefNet (togli sfondo)', + descrizione: 'Separa il soggetto dallo sfondo senza ridisegnarlo: i pixel restano quelli della tua foto, con capelli e bordi fini. Licenza MIT.', + file: '', cartella: 'segment/birefnet', url: '', byte: 444_473_596, circa: true, consigliato: false, + repo: BIREFNET, revisione: BIREFNET_REV, cerca: ['^model\\.safetensors$'], accanto: ['^config\\.json$'] + }, + { + id: 'seg-sam2', tipo: 'segmentatore', nome: 'SAM 2.1 small (estrai fedele)', + descrizione: 'Dal riquadro o dai punti che indichi ritaglia esattamente quella cosa dalla foto originale (persona, oggetto), senza ricrearla. Licenza Apache-2.0.', + file: '', cartella: 'segment/sam2.1-hiera-small', url: '', byte: 184_305_280, circa: true, consigliato: false, + repo: SAM2, revisione: SAM2_REV, cerca: ['^model\\.safetensors$'], accanto: ['^config\\.json$', '^preprocessor_config\\.json$', '^processor_config\\.json$'] + }, + // --- LLM uncensored per i testi (migliora prompt, descrivi foto, Idee, Estrai…) con Klein: stesso formato del Qwen3-VL 4B di Klein, quindi si sostituisce uno a uno --- + { + id: 'llm-huihui-4b', tipo: 'llm', uncensored: true, nome: 'Qwen3-VL 4B uncensored (Huihui, INT8)', + descrizione: 'Il Qwen3-VL 4B «abliterated» di huihui-ai, nel formato di ComfyUI. Scrive e descrive senza rifiuti: serve solo ai testi, non al disegno.', + file: '', cartella: 'text_encoders', url: '', byte: 5_064_440_625, circa: true, + repo: HUIHUI4B, revisione: HUIHUI4B_REV, cerca: ['^Huihui-Qwen3-VL-4B-Instruct-abliterated-int8_convrot\\.safetensors$', 'int8_convrot\\.safetensors$'] }, // --- VAE --- { @@ -178,4 +229,55 @@ export const CATALOGO: VoceCatalogo[] = [ } ] +// licenza, ruolo, «uncensored» e avvisi di ogni voce, in un posto solo (la UI li mostra accanto al nome; le licenze sono quelle dei repo, controllate il 22/10/2026: +// è una lettura tecnica dei testi, non un parere legale). Un componente «uncensored» ha il rifiuto tolto dai pesi: non cambia la licenza del modello che lo usa. +const AVVISO_QWEN = 'Qwen-Image 2.1 ha la licenza «Qwen Research»: solo ricerca e valutazione, niente uso commerciale senza un accordo con Qwen.' +const AVVISO_KOMMUNITY = 'Licenza FLUX Kommunity: l\'uso commerciale dei risultati ha condizioni (fatturato sotto 5 milioni di dollari, filtri o revisione dei risultati); leggila prima di pubblicare.' +const AVVISO_UNC = 'Questo componente non rifiuta nessuna richiesta e può produrre contenuti sensibili. Usalo solo per scopi leciti, con i diritti e il consenso necessari: sono vietati i contenuti sessuali con minori, le immagini intime non consensuali e ogni altro uso illegale. La licenza del modello resta valida e può limitare l\'uso commerciale.' + +const META: Record> = { + ...Object.fromEntries(['q4km', 'q6k', 'q8', 'turbo-q4', 'turbo-q6', 'turbo-q8'].map((id) => [id, { licenza: 'Qwen Research', licenzaUrl: LIC_QWEN, ruolo: 'denoiser', avviso: AVVISO_QWEN } as Partial])), + 'te-int8': { licenza: 'Apache-2.0', licenzaUrl: LIC_APACHE, ruolo: 'encoder' }, + ...Object.fromEntries(['te-ac-q4', 'te-ac-q6', 'te-ac-q8', 'te-mr-q4', 'te-mr-q6', 'te-mr-q8'].map((id) => [id, { licenza: 'Apache-2.0', licenzaUrl: LIC_APACHE, ruolo: 'encoder', uncensored: true, avviso: AVVISO_UNC } as Partial])), + 'trad-it-en': { licenza: 'CC-BY-4.0', licenzaUrl: 'https://creativecommons.org/licenses/by/4.0/' }, + ...Object.fromEntries(['klein-int8', 'klein-bf16', 'klein-base-int8', 'klein-base-bf16'].map((id) => [id, { licenza: 'FLUX Kommunity', licenzaUrl: LIC_KOMMUNITY, ruolo: 'denoiser', avviso: AVVISO_KOMMUNITY } as Partial])), + 'klein-te': { licenza: 'Apache-2.0', licenzaUrl: LIC_APACHE, ruolo: 'encoder' }, + 'klein-vae': { licenza: 'Apache-2.0', licenzaUrl: LIC_APACHE, ruolo: 'vae' }, + 'klein4b-fp8': { licenza: 'Apache-2.0', licenzaUrl: 'https://huggingface.co/black-forest-labs/FLUX.2-klein-4B/raw/main/LICENSE.md', ruolo: 'denoiser' }, + 'klein4b-te': { licenza: 'Apache-2.0', licenzaUrl: LIC_APACHE, ruolo: 'encoder' }, + 'sfondo-birefnet': { licenza: 'MIT', licenzaUrl: 'https://huggingface.co/ZhengPeng7/BiRefNet', ruolo: 'sfondo' }, + 'seg-sam2': { licenza: 'Apache-2.0', licenzaUrl: LIC_APACHE, ruolo: 'segmentatore' }, + 'llm-huihui-4b': { licenza: 'Apache-2.0', licenzaUrl: LIC_APACHE, ruolo: 'llm', avviso: AVVISO_UNC }, + 'lora-uc': { uncensored: true, licenza: 'Qwen Research', licenzaUrl: LIC_QWEN, ruolo: 'lora', avviso: AVVISO_QWEN + ' ' + AVVISO_UNC } +} + +export const CATALOGO: VoceCatalogo[] = GREZZO.map((v) => ({ ...v, ...(META[v.id] || {}) })) + export const voce = (id: string): VoceCatalogo | undefined => CATALOGO.find((v) => v.id === id) + +/** il testo di un avviso sui modelli uncensored, uguale per tutti (la UI lo mostra prima di chiedere l'accettazione) */ +export const AVVISO_UNCENSORED = AVVISO_UNC + +/** le licenze con cui un componente uncensored può stare nel catalogo: permissive, o «Qwen Research» (solo con l'avviso sul limite d'uso). + * Fuori: FLUX Kommunity (§4(d): niente scorciatoie per togliere le protezioni), qualunque «non-commercial» non letto, licenze mancanti o senza link. */ +export const LICENZE_UNCENSORED = ['Apache-2.0', 'MIT', 'BSD-3-Clause', 'CC-BY-4.0', 'Qwen Research'] + +/** una voce uncensored ha una licenza valida? null = sì, se no il motivo in italiano */ +export function motivoLicenza(v: VoceCatalogo): string | null { + if (!v.uncensored) return null + if (!v.licenza || !v.licenzaUrl) return `«${v.nome}» non ha una licenza dichiarata e verificabile: non si scarica.` + if (!LICENZE_UNCENSORED.includes(v.licenza)) return `«${v.nome}» ha la licenza «${v.licenza}», che non permette di distribuire una versione senza protezioni: non si scarica.` + if (v.licenza === 'Qwen Research' && !/Qwen Research/.test(v.avviso || '')) return `«${v.nome}» è sotto «Qwen Research» ma manca l'avviso sul limite d'uso: non si scarica.` + return null +} + +/** una voce uncensored si scarica solo con licenza valida (`motivoLicenza`), con l'interruttore «Modelli uncensored» acceso e l'avviso accettato; null = si può, se no il motivo in italiano. + * È la guardia del processo principale: la UI nasconde le voci, ma anche una chiamata diretta non le scarica. */ +export function motivoDiBlocco(v: VoceCatalogo, imp: { uncensored?: boolean; uncensoredAccettato?: boolean }): string | null { + if (!v.uncensored) return null + const licenza = motivoLicenza(v) + if (licenza) return licenza + if (!imp.uncensored) return `«${v.nome}» è un modello uncensored: accendi «Modelli uncensored» in Opzioni per poterlo scaricare.` + if (!imp.uncensoredAccettato) return `«${v.nome}» è un modello uncensored: prima leggi e accetta l'avviso in Opzioni → Modelli uncensored.` + return null +} diff --git a/src/main/installa/installatore.ts b/src/main/installa/installatore.ts index d79b679..497c9bd 100644 --- a/src/main/installa/installatore.ts +++ b/src/main/installa/installatore.ts @@ -8,7 +8,7 @@ import { totalmem } from 'node:os' import { basename, dirname, join, parse } from 'node:path' import { app } from 'electron' import type { ControlloSistema, PassoSetup, StatoVoce, VoceCatalogo } from '@shared/tipi' -import { CATALOGO, voce } from './catalogo' +import { CATALOGO, motivoDiBlocco, voce } from './catalogo' import { conRisolto, presenteRisolto, risolviVoce, risoltoSalvato, type FileRisolto } from './risolvi' import { elencoLora, scaricaLora } from '../lora' import { scarica } from './scarica' @@ -477,6 +477,9 @@ export async function installa(opz: OpzioniSetup, notifica: Notifica, provaMotor export async function scaricaVoce(id: string, avanz: (r: number, t: number, v: number) => void, segnale?: AbortSignal): Promise { const s = voce(id) if (!s) throw new Error('voce sconosciuta') + // i modelli uncensored si scaricano solo su richiesta esplicita: interruttore acceso e avviso accettato (la guardia sta qui, non solo nell'interfaccia) + const blocco = motivoDiBlocco(s, impostazioni()) + if (blocco) throw new Error(blocco) if (s.link) return scaricaLora(s.link, avanz, segnale) const imp = impostazioni() const intestazioni: Record = {} diff --git a/src/main/installa/risolvi.ts b/src/main/installa/risolvi.ts index c805a41..535f1f1 100644 --- a/src/main/installa/risolvi.ts +++ b/src/main/installa/risolvi.ts @@ -1,11 +1,12 @@ -// I modelli di terzi (Turbo di AtomicChat, encoder GGUF) hanno nomi di file che non conosciamo in anticipo: -// si chiede a Hugging Face l'elenco del repo e si sceglie il file giusto (Q4 / Q6 / Q8) con le espressioni del catalogo. +// I modelli di terzi (Turbo di AtomicChat, encoder GGUF, Klein, segmentatori) hanno nomi di file che non conosciamo in anticipo: +// si chiede a Hugging Face l'elenco del repo (alla revisione fissata nel catalogo, se c'è) e si sceglie il file giusto con le espressioni del catalogo. // La scelta si ricorda in `risolti.json`, così si sa cosa c'è sul disco anche senza rete. -import { existsSync, readFileSync, statSync, writeFileSync } from 'node:fs' +import { existsSync, readFileSync, statSync } from 'node:fs' import { join } from 'node:path' import type { VoceCatalogo } from '@shared/tipi' import { DATI, scriviAtomico } from '../percorsi' import { impostazioni } from '../impostazioni' +import { TIPI_SAFETENSORS, scegli, tronco, type VoceHF } from './scegli' export interface FileRisolto { file: string @@ -17,7 +18,7 @@ export interface FileRisolto { export interface Risolto extends FileRisolto { /** l'mmproj del Qwen3-VL (serve alla vista): sta nella stessa cartella, col nome che il nodo GGUF si aspetta */ mmproj?: FileRisolto - /** gli altri file accanto (config e tokenizer del traduttore) */ + /** gli altri file accanto (config e tokenizer del traduttore, config del segmentatore) */ altri?: FileRisolto[] } @@ -43,38 +44,18 @@ export function senzaQuant(nome: string): string { return base.replace(/[-_]?(?:ud-)?i?q[0-9]_[a-z0-9_-]{1,8}$/i, '') } -interface VoceHF { - type: string - path: string - size?: number - lfs?: { oid?: string; size?: number } -} - -const tronco = (p: string): string => p.split('/').pop() || p - -async function elenco(repo: string): Promise { +async function elenco(repo: string, revisione?: string): Promise { const imp = impostazioni() const intestazioni: Record = {} if (imp.tokenHF) intestazioni.Authorization = 'Bearer ' + imp.tokenHF - const r = await fetch(`https://huggingface.co/api/models/${repo}/tree/main?recursive=true`, { headers: intestazioni }) + const r = await fetch(`https://huggingface.co/api/models/${repo}/tree/${revisione || 'main'}?recursive=true`, { headers: intestazioni }) if (!r.ok) throw new Error(`Hugging Face non risponde per ${repo} (${r.status})`) return ((await r.json()) as VoceHF[]).filter((f) => f.type === 'file') } -function scegli(files: VoceHF[], cerca: string[], escludiMmproj: boolean, ext = /\.gguf$/i): VoceHF | undefined { - const candidati = files.filter((f) => ext.test(f.path) && (/mmproj/i.test(f.path) ? !escludiMmproj : escludiMmproj)) - for (const re of cerca) { - const r = new RegExp(re, 'i') - // si prova col nome del file e poi col percorso intero (alcuni repo mettono il tipo nella cartella: int8/…) - const trovato = candidati.find((f) => r.test(tronco(f.path))) || candidati.find((f) => r.test(f.path)) - if (trovato) return trovato - } - return undefined -} - -const aFile = (repo: string, f: VoceHF, nome: string): FileRisolto => ({ +const aFile = (repo: string, f: VoceHF, nome: string, revisione?: string): FileRisolto => ({ file: nome, - url: `https://huggingface.co/${repo}/resolve/main/${f.path}`, + url: `https://huggingface.co/${repo}/resolve/${revisione || 'main'}/${f.path}`, byte: f.lfs?.size ?? f.size ?? 0, sha256: f.lfs?.oid }) @@ -82,16 +63,16 @@ const aFile = (repo: string, f: VoceHF, nome: string): FileRisolto => ({ /** sceglie i file di una voce del catalogo (con `repo`) e se li ricorda; serve la rete */ export async function risolviVoce(v: VoceCatalogo): Promise { if (!v.repo || !v.cerca) throw new Error(`${v.nome}: non è una voce da cercare`) - const files = await elenco(v.repo) - // i pesi sono un .gguf (denoiser, encoder) o, per il traduttore, un .safetensors / .bin - const pesi = v.tipo === 'traduttore' || v.tipo === 'klein' || v.tipo === 'klein-te' || v.tipo === 'klein-vae' ? /\.(safetensors|bin)$/i : /\.gguf$/i + const files = await elenco(v.repo, v.revisione) + // i pesi sono un .gguf (denoiser, encoder) o, per il traduttore, Klein, il segmentatore…, un .safetensors / .bin + const pesi = TIPI_SAFETENSORS.has(v.tipo) ? /\.(safetensors|bin)$/i : /\.gguf$/i const principale = scegli(files, v.cerca, true, pesi) if (!principale) { const elencoFile = files.filter((f) => pesi.test(f.path)).map((f) => f.path).slice(0, 40).join(', ') throw new Error(`Nel repo ${v.repo} non trovo il file giusto per «${v.nome}». File del repo: ${elencoFile || '(nessuno)'}`) } const nome = tronco(principale.path) - const r: Risolto = aFile(v.repo, principale, nome) + const r: Risolto = aFile(v.repo, principale, nome, v.revisione) if (v.cercaMmproj) { const m = scegli(await elenco(v.mmprojRepo || v.repo), v.cercaMmproj, false) if (!m) throw new Error(`Nel repo ${v.mmprojRepo || v.repo} non trovo l'mmproj (la parte che fa guardare le foto)`) @@ -102,7 +83,7 @@ export async function risolviVoce(v: VoceCatalogo): Promise { if (v.accanto) { r.altri = v.accanto.flatMap((re) => { const f = files.find((x) => new RegExp(re, 'i').test(tronco(x.path))) - return f ? [aFile(v.repo!, f, tronco(f.path))] : [] + return f ? [aFile(v.repo!, f, tronco(f.path), v.revisione)] : [] }) } ricorda(v.id, r) diff --git a/src/main/installa/scegli.ts b/src/main/installa/scegli.ts new file mode 100644 index 0000000..5d1e77a --- /dev/null +++ b/src/main/installa/scegli.ts @@ -0,0 +1,24 @@ +// Sceglie il file giusto fra quelli di un repo Hugging Face con le espressioni del catalogo (in ordine di preferenza). +// File puro (niente Electron): lo prova `node test/catalogo.mjs` con gli alberi veri dei repo. +export interface VoceHF { + type: string + path: string + size?: number + lfs?: { oid?: string; size?: number } +} + +export const tronco = (p: string): string => p.split('/').pop() || p + +export function scegli(files: VoceHF[], cerca: string[], escludiMmproj: boolean, ext = /\.gguf$/i): VoceHF | undefined { + const candidati = files.filter((f) => ext.test(f.path) && (/mmproj/i.test(f.path) ? !escludiMmproj : escludiMmproj)) + for (const re of cerca) { + const r = new RegExp(re, 'i') + // si prova col nome del file e poi col percorso intero (alcuni repo mettono il tipo nella cartella: int8/…) + const trovato = candidati.find((f) => r.test(tronco(f.path))) || candidati.find((f) => r.test(f.path)) + if (trovato) return trovato + } + return undefined +} + +/** i tipi di voce i cui pesi sono .safetensors/.bin (gli altri, denoiser e encoder GGUF, sono .gguf) */ +export const TIPI_SAFETENSORS = new Set(['traduttore', 'klein', 'klein-te', 'klein-vae', 'sfondo', 'segmentatore', 'llm']) diff --git a/src/main/lora.ts b/src/main/lora.ts index 65f4964..77a54c3 100644 --- a/src/main/lora.ts +++ b/src/main/lora.ts @@ -1,10 +1,11 @@ // I LoRA: i file stanno in models/loras, le note (nome, parole chiave, forza, anteprima) in lora.json. // Si importano trascinandoli, si scaricano da un link (Hugging Face, Civitai o diretto), -// e dall'intestazione safetensors si capisce se sono per Qwen-Image 2.1. +// e dall'intestazione safetensors si capisce per quale modello sono (Qwen-Image 2.1, FLUX 3 Klein, FLUX.2 Klein 4B: `shared/loraBase.ts`). import { shell } from 'electron' import { copyFileSync, existsSync, openSync, readFileSync, readSync, closeSync, readdirSync, statSync, writeFileSync } from 'node:fs' import { basename, extname, join, relative, resolve, sep } from 'node:path' -import type { InfoLora } from '@shared/tipi' +import type { Famiglia, InfoLora } from '@shared/tipi' +import { analizzaChiavi } from '@shared/loraBase' import { DATI, assicura, scriviAtomico } from './percorsi' import { impostazioni } from './impostazioni' import { scarica } from './installa/scarica' @@ -48,27 +49,15 @@ function intestazione(percorso: string): Record | null { } /** Qwen-Image 2.1 è a flusso singolo (solo img_mlp); Qwen-Image 1.x ha txt_mlp/add_q_proj; SD/SDXL/Flux hanno altri nomi */ -function analizza(percorso: string): { compatibile: InfoLora['compatibile']; base?: string; parole?: string; nome?: string } { +function analizza(percorso: string): { compatibile: InfoLora['compatibile']; base?: string; famiglie?: Famiglia[]; parole?: string; nome?: string } { const h = intestazione(percorso) if (!h) return { compatibile: '?' } const chiavi = Object.keys(h).filter((k) => k !== '__metadata__') const meta = (h.__metadata__ || {}) as Record - const tutte = chiavi.join('\n') - let compatibile: InfoLora['compatibile'] = '?' - let base: string | undefined - if (/txt_mlp|add_q_proj|txt_mod|img_mod\./.test(tutte)) { - compatibile = 'no' - base = 'Qwen-Image 1.x' - } else if (/lora_unet_|down_blocks|input_blocks|lora_te/.test(tutte)) { - compatibile = 'no' - base = 'Stable Diffusion / SDXL' - } else if (/double_blocks|single_blocks/.test(tutte)) { - compatibile = 'no' - base = 'Flux' - } else if (/transformer_blocks\.\d+\.(attn\.to_q|img_mlp)/.test(tutte)) { - compatibile = 'si' - base = 'Qwen-Image 2.1' - } + const forme: Record = {} + for (const k of chiavi) forme[k] = (h[k] as { shape?: number[] } | undefined)?.shape + const { compatibile, famiglie, base: baseRiconosciuta } = analizzaChiavi(chiavi, forme) + let base = baseRiconosciuta const archi = meta['modelspec.architecture'] || meta.ss_base_model_version || '' if (archi) base = base ? `${base} (${archi})` : archi let parole = meta['modelspec.trigger_phrase'] || '' @@ -86,7 +75,7 @@ function analizza(percorso: string): { compatibile: InfoLora['compatibile']; bas /* niente */ } } - return { compatibile, base, parole, nome: meta['modelspec.title'] || meta.ss_output_name } + return { compatibile, base, famiglie, parole, nome: meta['modelspec.title'] || meta.ss_output_name } } export function elencoLora(): InfoLora[] { @@ -112,10 +101,10 @@ export function elencoLora(): InfoLora[] { // ComfyUI vuole il nome relativo alla cartella loras, con le barre di Windows const file = relative(c, p) let nota = n[file] - if (!nota || nota.compatibile === undefined) { + if (!nota || nota.compatibile === undefined || nota.famiglie === undefined) { const a = analizza(p) const delCatalogo = CATALOGO.find((x) => x.tipo === 'lora' && x.file === v) - nota = { nome: delCatalogo?.nome || a.nome || basename(v, extname(v)), parole: a.parole || '', forza: 1, aggiunta: st.mtimeMs, ...nota, compatibile: a.compatibile, base: a.base } + nota = { nome: delCatalogo?.nome || a.nome || basename(v, extname(v)), parole: a.parole || '', forza: 1, aggiunta: st.mtimeMs, ...nota, compatibile: a.compatibile, base: a.base, famiglie: a.famiglie } n[file] = nota cambiate = true } diff --git a/src/main/motore/adatta.ts b/src/main/motore/adatta.ts new file mode 100644 index 0000000..0f80538 --- /dev/null +++ b/src/main/motore/adatta.ts @@ -0,0 +1,28 @@ +// Dalla richiesta che manda la GUI v2 a quella che capisce il motore. File puro (niente Electron): lo prova `node test/adatta.mjs`. +// +// - `box` (Editor box: BoxRegione con rect/tipo/rigidita/ordine) → `regioni` (riquadro + prompt, nell'ordine dell'editor; le nascoste non partono). +// Se ci sono riquadri validi, il prompt della scena è `promptScena` (se c'è); senza riquadri il prompt resta com'è. +// - `sfondo` → 'trasparente' = `trasparente: true`; { tinta } = una riga in più nel prompt («on a plain solid #hex background»); 'scena' = niente. +// - la rigidità `blocca` (maschera e compositing) per ora vale come `decisa`; `rif` (foto di riferimento del box) non è ancora usato dal motore. +import type { Regione, Richiesta } from '../../shared/tipi' + +const HEX = /^#?[0-9a-f]{3}([0-9a-f]{3})?$/i + +export function adattaRichiesta(q: Richiesta): Richiesta { + let fuori = q + if (q.box && q.box.length) { + const visibili = q.box.filter((b) => !b.nascosta && (b.prompt || '').trim() && Array.isArray(b.rect) && b.rect.length === 4) + const regioni: Regione[] = [...visibili].sort((a, b) => (a.ordine ?? 0) - (b.ordine ?? 0)).map((b) => ({ id: b.id, box: b.rect, prompt: b.prompt })) + if (regioni.length) { + const scena = (q.promptScena ?? '').trim() + fuori = { ...fuori, regioni: [...(q.regioni || []), ...regioni], prompt: scena || fuori.prompt } + } + } + const s = q.sfondo + if (s === 'trasparente') fuori = { ...fuori, trasparente: true } + else if (s && typeof s === 'object' && typeof s.tinta === 'string' && HEX.test(s.tinta.trim())) { + const colore = s.tinta.trim().startsWith('#') ? s.tinta.trim() : '#' + s.tinta.trim() + fuori = { ...fuori, prompt: `${fuori.prompt.trim().replace(/[.\s]+$/, '')}, on a plain solid ${colore} background` } + } + return fuori +} diff --git a/src/main/motore/grafi.ts b/src/main/motore/grafi.ts index ae0591d..1ff0ccb 100644 --- a/src/main/motore/grafi.ts +++ b/src/main/motore/grafi.ts @@ -2,8 +2,9 @@ // modello GGUF → (Turbo) → LoRA → cache KV di Qwen-Image 2.1 · text encoder Qwen3-VL · VAE // e i parametri della pipeline ufficiale (euler/simple, cfg 1, 40 passi; 8 col Turbo). import { modelloTurbo } from '@shared/tipi' -import type { Famiglia, Impostazioni, Richiesta, Riquadro } from '@shared/tipi' +import type { Famiglia, Impostazioni, LoraAttiva, Richiesta, Riquadro } from '@shared/tipi' import type { ChiaviRtx } from './cliente' +import { promptConPosizioni, promptKleinConBox, regioniPronte } from './regioni' /** il profilo di memoria: con poca VRAM il VAE lavora a tessere (niente picchi che finiscono nella RAM condivisa) */ export type Memoria = 'bassa' | 'normale' | 'alta' @@ -79,18 +80,29 @@ function caricaClipPerTesto(g: Grafo, imp: Impostazioni, kleinEncoder: boolean | /** il modello che disegna questa richiesta */ export const famigliaDi = (q: Richiesta, imp: Impostazioni): Famiglia => q.famiglia ?? imp.famiglia ?? 'qwen' -/** FLUX 3 Klein: UNETLoader (i pesi INT8/BF16 li riconosce da soli ComfyUI), text encoder Qwen3-VL 4B di tipo flux3, VAE di Klein. Niente LoRA né cache KV. */ -function baseKlein(g: Grafo, imp: Impostazioni, fam: Famiglia): { model: Rif; clip: Rif; vae: Rif } { +/** i LoRA accesi in fila (l'ordine della lista è l'ordine di applicazione): un LoraLoaderModelOnly per ciascuno, forza ≠ 0. + * Vale per tutte le famiglie: Qwen e le Klein. `salta` = file da non mettere (il Turbo, già messo a parte). */ +function applicaLora(g: Grafo, model: Rif, lora: LoraAttiva[], salta?: string): Rif { + let m = model + for (const l of lora) { + if (!l.file || !l.forza || l.file === salta) continue + m = r(g.add('LoraLoaderModelOnly', { model: m, lora_name: l.file, strength_model: Math.max(-2, Math.min(2, l.forza)) }, 'Applico i LoRA')) + } + return m +} + +/** FLUX 3 Klein e FLUX.2 Klein 4B: UNETLoader (i pesi INT8/FP8/BF16 li riconosce da soli ComfyUI), LoRA, text encoder (Qwen3-VL 4B di tipo flux3; Qwen3 4B di tipo flux2 per la 4B), VAE di FLUX.2. Niente cache KV. */ +function baseKlein(g: Grafo, imp: Impostazioni, fam: Famiglia, lora: LoraAttiva[]): { model: Rif; clip: Rif; vae: Rif } { const file = fam === 'klein-base' ? imp.kleinBase : imp.kleinModello - const model = r(g.add('UNETLoader', { unet_name: file, weight_dtype: 'default' }, 'Carico il modello')) - const clip = r(g.add('CLIPLoader', { clip_name: imp.kleinEncoder, type: 'flux3', device: imp.encoderSuCpu ? 'cpu' : 'default' }, 'Carico il text encoder')) + const model = applicaLora(g, r(g.add('UNETLoader', { unet_name: file, weight_dtype: 'default' }, 'Carico il modello')), lora) + const clip = r(g.add('CLIPLoader', { clip_name: imp.kleinEncoder, type: fam === 'klein4b' ? 'flux2' : 'flux3', device: imp.encoderSuCpu ? 'cpu' : 'default' }, 'Carico il text encoder')) const vae = r(g.add('VAELoader', { vae_name: imp.kleinVae }, 'Carico il VAE')) return { model, clip, vae } } function base(g: Grafo, q: Richiesta, imp: Impostazioni): { model: Rif; clip: Rif; vae: Rif } { const fam = famigliaDi(q, imp) - if (fam !== 'qwen') return baseKlein(g, imp, fam) + if (fam !== 'qwen') return baseKlein(g, imp, fam, q.lora) const carico = 'Carico il modello' const unet = imp.modello.toLowerCase().endsWith('.gguf') ? g.add('UnetLoaderGGUF', { unet_name: imp.modello }, carico) @@ -99,10 +111,7 @@ function base(g: Grafo, q: Richiesta, imp: Impostazioni): { model: Rif; clip: Ri const turbo = conTurbo(q, imp) const nativo = modelloTurbo(imp.modello) if (turbo && !nativo) model = r(g.add('LoraLoaderModelOnly', { model, lora_name: imp.loraTurbo, strength_model: 1 }, 'Accendo il Turbo')) - for (const l of q.lora) { - if (!l.file || !l.forza || (turbo && l.file === imp.loraTurbo)) continue - model = r(g.add('LoraLoaderModelOnly', { model, lora_name: l.file, strength_model: l.forza }, 'Applico i LoRA')) - } + model = applicaLora(g, model, q.lora, turbo ? imp.loraTurbo : undefined) model = r(g.add('QwenImage21Cache', { model, device: imp.cacheKV, dtype: imp.cacheTipo })) const clip = caricaClip(g, imp) const vae = r(g.add('VAELoader', { vae_name: imp.vae }, 'Carico il VAE')) @@ -166,12 +175,35 @@ function condiziona( return { pos, neg } } +/** il testo di Crea o Modifica con i riquadri dell'Editor box: Klein 3 li legge (JSON con le coordinate), gli altri modelli li hanno a parole. Senza riquadri valide, il testo com'è. */ +function conRegioni(testo: string, q: Richiesta, fam: Famiglia, modifica: boolean): string { + const regs = regioniPronte(q.regioni) + if (!regs.length) return testo + return fam === 'klein' || fam === 'klein-base' ? promptKleinConBox(testo, regs, modifica) : promptConPosizioni(testo, regs, modifica) +} + /** la latente vuota per disegnare da zero a quella misura */ function latenteVuota(g: Grafo, fam: Famiglia, w: number, h: number): Rif { return r(g.add(fam === 'qwen' ? 'EmptyLatentImage' : 'EmptyFlux2LatentImage', { width: w, height: h, batch_size: 1 })) } -function campiona(g: Grafo, q: Richiesta, fam: Famiglia, model: Rif, c: Condizioni, latente: Rif, passi: number, denoise = 1): Rif { +function campiona(g: Grafo, q: Richiesta, fam: Famiglia, model: Rif, c: Condizioni, latente: Rif, passi: number, denoise = 1, dim?: { w: number; h: number }): Rif { + if (fam === 'klein4b') { + // FLUX.2 Klein 4B (distillato): come il workflow ufficiale di ComfyUI, scaletta Flux2Scheduler (dipende dalla grandezza) + CFGGuider 1 + euler. + // Con denoise < 1 la scaletta è più lunga e se ne prende la coda (SplitSigmasDenoise): si fanno `passi` passi veri. + const w = dim?.w ?? q.larghezza + const h = dim?.h ?? q.altezza + const totali = denoise < 0.999 ? Math.max(passi, Math.round(passi / Math.max(0.05, denoise))) : passi + const scaletta = g.add('Flux2Scheduler', { steps: totali, width: w, height: h }) + const sigme: Rif = denoise < 0.999 ? r(g.add('SplitSigmasDenoise', { sigmas: r(scaletta), denoise }), 1) : r(scaletta) + const guida = g.add('CFGGuider', { model, positive: c.pos, negative: c.neg, cfg: 1 }) + const out = g.add( + 'SamplerCustomAdvanced', + { noise: r(g.add('RandomNoise', { noise_seed: q.seed })), guider: r(guida), sampler: r(g.add('KSamplerSelect', { sampler_name: 'euler' })), sigmas: sigme, latent_image: latente }, + 'Disegno' + ) + return r(out, 0) + } // Klein veloce è distillato (8 passi, senza CFG); Klein Base vuole CFG 3,5; sampler e scheduler sono quelli del riferimento (euler, e la scaletta del modello) const cfg = fam === 'qwen' ? q.cfg : fam === 'klein' ? 1 : q.cfg > 1 ? q.cfg : 3.5 return r( @@ -201,9 +233,11 @@ function aTessere(g: Grafo, mp: number): boolean { const TESSERE = { tile_size: 512, overlap: 64, temporal_size: 64, temporal_overlap: 8 } /** decodifica e (se non trasparente) toglie il canale alfa: il VAE di Qwen 2.1 esce sempre RGBA */ -function sviluppa(g: Grafo, q: Richiesta, lat: Rif, vae: Rif, mp: number): Rif { +function sviluppa(g: Grafo, q: Richiesta, lat: Rif, vae: Rif, mp: number, togliSfondo = false): Rif { const fase = 'Sviluppo l\'immagine' const dec = aTessere(g, mp) ? g.add('VAEDecodeTiled', { samples: lat, vae, ...TESSERE }, fase) : g.add('VAEDecode', { samples: lat, vae }, fase) + // i modelli Klein non disegnano la trasparenza (Qwen sì, col suo RGBA): «senza sfondo» = il soggetto si ritaglia con BiRefNet, i pixel restano quelli disegnati + if (q.trasparente && togliSfondo) return r(g.add('DaProdSfondo', { image: r(dec), modello: 'birefnet', pulisci_bordi: true, ritaglia: false }, 'Tolgo lo sfondo'), 0) if (q.trasparente) return r(dec) return r(g.add('SplitImageWithAlpha', { image: r(dec) }), 0) } @@ -343,6 +377,29 @@ export function costruisci(q: Richiesta, imp: Impostazioni, ing: Ingressi, prefi return { prompt: g.nodi, uscita, fasi: g.fasi, megapixel: 0, passi: q.maxToken || 400 } } + // togli sfondo: la foto originale più il canale alfa di BiRefNet (nessun modello che disegna: niente VRAM per il DiT) + if (q.modalita === 'sfondo') { + const img = r(g.add('LoadImage', { image: ing.immagini[0].nome }, 'Carico la foto')) + const fuori = r(g.add('DaProdSfondo', { image: img, modello: 'birefnet', pulisci_bordi: true, ritaglia: false }, 'Tolgo lo sfondo'), 0) + return { prompt: g.nodi, uscita: salva(g, fuori, prefisso), fasi: g.fasi, megapixel: 0, passi: 1 } + } + + // estrai fedele: SAM 2 ritaglia dai pixel originali la cosa indicata (riquadro e/o punti, 0–1 sulla foto) + if (q.modalita === 'segmenta') { + const s = q.segmenta + if (!s || (!s.box && !(s.punti && s.punti.length))) throw new Error('dimmi cosa estrarre: disegna un riquadro o tocca dei punti') + const img = r(g.add('LoadImage', { image: ing.immagini[0].nome }, 'Carico la foto')) + const fuori = r( + g.add( + 'DaProdSegmenta', + { image: img, modello: 'sam2.1-hiera-small', box: s.box ? JSON.stringify(s.box) : '', punti: s.punti && s.punti.length ? JSON.stringify(s.punti) : '', ritaglia: s.ritaglia !== false, margine: 12, pulisci_bordi: true }, + 'Ritaglio la cosa scelta' + ), + 0 + ) + return { prompt: g.nodi, uscita: salva(g, fuori, prefisso), fasi: g.fasi, megapixel: 0, passi: 1 } + } + if (q.modalita === 'ingrandisci') { const img = r(g.add('LoadImage', { image: ing.immagini[0].nome }, 'Carico la foto')) const f = q.fattore || 2 @@ -359,7 +416,8 @@ export function costruisci(q: Richiesta, imp: Impostazioni, ing: Ingressi, prefi const leggo = 'Leggo il prompt e le foto' if (q.modalita === 'crea') { - const testo = q.trasparente && fam === 'qwen' ? AVVOLGI_TRASPARENTE(q.prompt) : q.prompt + let testo = q.trasparente && fam === 'qwen' ? AVVOLGI_TRASPARENTE(q.prompt) : q.prompt + testo = conRegioni(testo, q, fam, false) // con dei riferimenti (…) il modello li guarda mentre disegna da zero: persone, vestiti, oggetti, stili const caricate = ing.immagini.slice(0, fam === 'qwen' ? 16 : 10).map((im) => r(g.add('LoadImage', { image: im.nome }, 'Carico i riferimenti'))) const cond = condiziona(g, fam, { clip, vae }, { prompt: testo, neg, resolution: 1024, conVae: caricate.length > 0, immagini: caricate, mpRiferimenti: 1, scalaDa: 0, fase: leggo }) @@ -368,9 +426,9 @@ export function costruisci(q: Richiesta, imp: Impostazioni, ing: Ingressi, prefi const m = fam === 'qwen' && conTurbo(q, imp) && !modelloTurbo(imp.modello) ? r(g.add('ModelSamplingFlux', { model, max_shift: 0.6935, base_shift: 0.5, width: q.larghezza, height: q.altezza })) : model - const ks = campiona(g, q, fam, m, cond, lat, q.passi) + const ks = campiona(g, q, fam, m, cond, lat, q.passi, 1, { w: q.larghezza, h: q.altezza }) const mp = (q.larghezza * q.altezza) / 1048576 - const uscita = salva(g, allaMisura(g, q, imp, sviluppa(g, q, ks, vae, mp), q.larghezza, q.altezza, rtx), prefisso) + const uscita = salva(g, allaMisura(g, q, imp, sviluppa(g, q, ks, vae, mp, fam !== 'qwen'), q.larghezza, q.altezza, rtx), prefisso) return { prompt: g.nodi, uscita, fasi: g.fasi, megapixel: mp, passi: q.passi } } @@ -379,11 +437,11 @@ export function costruisci(q: Richiesta, imp: Impostazioni, ing: Ingressi, prefi const res = Math.round((Math.sqrt(q.megapixel) * 1024) / 32) * 32 const d = dimensioniPer(q.megapixel, ing.immagini[0].larghezza, ing.immagini[0].altezza) const caricate = ing.immagini.slice(0, fam === 'qwen' ? 16 : 10).map((im) => r(g.add('LoadImage', { image: im.nome }, 'Carico le foto'))) - const cond = condiziona(g, fam, { clip, vae }, { prompt: q.prompt, neg, resolution: res, conVae: true, immagini: caricate, misuraPrima: d, fase: leggo }) + const cond = condiziona(g, fam, { clip, vae }, { prompt: conRegioni(q.prompt, q, fam, true), neg, resolution: res, conVae: true, immagini: caricate, misuraPrima: d, fase: leggo }) const lat = cond.lat ?? latenteVuota(g, fam, d.w, d.h) - const ks = campiona(g, q, fam, model, cond, lat, q.passi) + const ks = campiona(g, q, fam, model, cond, lat, q.passi, 1, d) const mp = (d.w * d.h) / 1048576 - const uscita = salva(g, allaMisura(g, q, imp, sviluppa(g, q, ks, vae, mp), d.w, d.h, rtx), prefisso) + const uscita = salva(g, allaMisura(g, q, imp, sviluppa(g, q, ks, vae, mp, fam !== 'qwen'), d.w, d.h, rtx), prefisso) const area = { x: 0, y: 0, w: ing.immagini[0].larghezza, h: ing.immagini[0].altezza } return { prompt: g.nodi, uscita, fasi: g.fasi, megapixel: mp, passi: q.passi, area } } @@ -412,8 +470,8 @@ export function costruisci(q: Richiesta, imp: Impostazioni, ing: Ingressi, prefi const lat = codifica(g, img, vae, mp) const forza = q.forza ?? (q.modalita === 'rifinisci' ? 0.45 : 0.65) const passi = passiPer(q.passi, forza, minimo) - const ks = campiona(g, q, fam, model, cond, lat, passi, forza) - const uscita = salva(g, allaMisura(g, q, imp, sviluppa(g, q, ks, vae, mp), W, H, rtx), prefisso) + const ks = campiona(g, q, fam, model, cond, lat, passi, forza, { w: W, h: H }) + const uscita = salva(g, allaMisura(g, q, imp, sviluppa(g, q, ks, vae, mp, fam !== 'qwen'), W, H, rtx), prefisso) return { prompt: g.nodi, uscita, fasi: g.fasi, megapixel: mp, passi } } @@ -525,7 +583,7 @@ export function costruisci(q: Richiesta, imp: Impostazioni, ing: Ingressi, prefi const dd = r(g.add('DifferentialDiffusion', { model, strength: 1 })) const forza = q.forza ?? 1 const passi = passiPer(q.passi, forza, minimo) - const ks = campiona(g, q, fam, dd, cond, lat, passi, forza) + const ks = campiona(g, q, fam, dd, cond, lat, passi, forza, { w: box.tw, h: box.th }) // il doppio passaggio nel VAE sposta un filo i colori: si misurano appena fuori dalla zona e si rimettono a posto const fatto = r(g.add('DaProdAccordaColori', { image: sviluppa(g, { ...q, trasparente: false }, ks, vae, mp), reference: grande, mask: mDura, strength: 1 })) const indietro = r(g.add('ImageScale', { image: fatto, upscale_method: 'lanczos', width: box.w, height: box.h, crop: 'disabled' }, 'Incollo')) diff --git a/src/main/motore/klein.ts b/src/main/motore/klein.ts index c980218..000ec4d 100644 --- a/src/main/motore/klein.ts +++ b/src/main/motore/klein.ts @@ -1,4 +1,4 @@ -// FLUX 3 Klein: quali file usare. Le impostazioni possono dirlo; se sono vuote si usano quelli scaricati dall'app (risolti.json) e, +// FLUX 3 Klein e FLUX.2 Klein 4B: quali file usare. Le impostazioni possono dirlo; se sono vuote si usano quelli scaricati dall'app (risolti.json) e, // in mancanza, si cerca nelle cartelle dei modelli (chi ha già i file li mette lì a mano). import { existsSync, readdirSync } from 'node:fs' import { join } from 'node:path' @@ -28,30 +28,66 @@ const preferisci = (nomi: string[], ordine: RegExp[]): string => { return '' } -const NON_PESI = /vae|qwen|text_enc|mmproj/i +const NON_PESI = /vae|qwen|text_enc|mmproj|upsampl/i +/** i file di FLUX.2 Klein 4B / 9B (altra architettura di FLUX 3): non vanno mai presi per quelli di Klein */ +const E_FLUX2 = /klein.*[-_](4b|9b)|[-_](4b|9b).*klein/i -/** il Qwen3-VL 4B di Klein (impostazioni → scaricato dall'app → cercato nelle cartelle), o '' se non c'è: serve anche a scrivere testo senza il modello che disegna */ +/** il Qwen3-VL 4B di Klein (impostazioni → scaricato dall'app → cercato nelle cartelle), o '' se non c'è: serve anche a scrivere testo senza il modello che disegna. + * Non prende mai il Qwen3 4B di solo testo di FLUX.2 Klein 4B (`qwen_3_4b.safetensors`: è un altro encoder) né l'LLM uncensored (si sceglie a parte: `encoderDiTesto`). */ export function encoderKleinFile(imp: Impostazioni): string { const encoder = sul(imp, 'text_encoders', imp.kleinEncoder) || sul(imp, 'text_encoders', risoltoSalvato('klein-te')?.file) if (encoder) return encoder - return preferisci(lista(imp, 'text_encoders'), [/^qwen3vl_4b_int8/i, /(qwen3|qwen_3).*vl.*4b.*(int8|fp8|q8)/i, /(qwen3|qwen_3).*4b.*(int8|fp8|q8)/i, /(qwen3|qwen_3).*vl.*4b/i, /(qwen3|qwen_3).*4b/i, /klein/i]) + const candidati = lista(imp, 'text_encoders').filter((n) => !/^qwen_3_4b/i.test(n) && !/abliterated|uncensored/i.test(n)) + return preferisci(candidati, [/^qwen3vl_4b_int8/i, /(qwen3|qwen_3).*vl.*4b.*(int8|fp8|q8)/i, /(qwen3|qwen_3).*4b.*(int8|fp8|q8)/i, /(qwen3|qwen_3).*vl.*4b/i, /klein/i]) } -/** i file di Klein (o di Klein Base) da usare adesso, o l'errore da mostrare in italiano se mancano */ +/** l'LLM uncensored per i testi (voce `llm-huihui-4b`), se l'interruttore è acceso, l'avviso accettato, la scelta fatta e il file c'è; se no '' */ +export function llmUncensoredFile(imp: Impostazioni): string { + if (!imp.uncensored || !imp.uncensoredAccettato || !imp.testoUncensored) return '' + return sul(imp, 'text_encoders', risoltoSalvato('llm-huihui-4b')?.file) +} + +/** chi scrive i testi con le famiglie Klein: l'LLM uncensored se scelto, se no il Qwen3-VL 4B di Klein ('' = nessuno dei due: si resta con l'encoder di Qwen). + * L'LLM uncensored scrive solo i testi: il conditioning del disegno resta sempre l'encoder ufficiale (`fileKlein`). */ +export function encoderDiTesto(imp: Impostazioni): string { + return llmUncensoredFile(imp) || encoderKleinFile(imp) +} + +/** i file di Klein (o di Klein Base, o di FLUX.2 Klein 4B) da usare adesso, o l'errore da mostrare in italiano se mancano */ export function fileKlein(imp: Impostazioni, famiglia: Famiglia): FileKlein | string { + if (famiglia === 'klein4b') return fileKlein4b(imp) const base = famiglia === 'klein-base' const idModelli = base ? ['klein-base-int8', 'klein-base-bf16'] : ['klein-int8', 'klein-bf16'] const scelto = base ? imp.kleinBase : imp.kleinModello let modello = sul(imp, 'diffusion_models', scelto) if (!modello) for (const id of idModelli) modello = modello || sul(imp, 'diffusion_models', risoltoSalvato(id)?.file) if (!modello) { - const tutti = [...lista(imp, 'diffusion_models'), ...lista(imp, 'unet')].filter((n) => /klein/i.test(n) && !NON_PESI.test(n) && /base/i.test(n) === base) + const tutti = [...lista(imp, 'diffusion_models'), ...lista(imp, 'unet')].filter((n) => /klein/i.test(n) && !NON_PESI.test(n) && !E_FLUX2.test(n) && /base/i.test(n) === base) modello = preferisci(tutti, [/int8/i, /bf16/i, /fp8/i, /./]) } const encoder = encoderKleinFile(imp) - let vae = sul(imp, 'vae', imp.kleinVae) || sul(imp, 'vae', risoltoSalvato('klein-vae')?.file) - if (!vae) vae = preferisci(lista(imp, 'vae'), [/^flux2-vae\./i, /klein/i, /flux.?[23]/i, /^ae\./i, /flux/i]) + const vae = vaeKlein(imp) const mancano = [!modello && (base ? 'il modello Klein Base' : 'il modello Klein'), !encoder && 'il text encoder (Qwen3-VL 4B)', !vae && 'il VAE'].filter(Boolean) if (mancano.length) return `FLUX 3 Klein non è pronto: mancano ${mancano.join(', ')}. Scaricali dal selettore del modello (o da Opzioni → Scarica altri modelli).` return { modello, encoder, vae } } + +/** il VAE di FLUX.2 (lo stesso per Klein 3 e Klein 4B) */ +function vaeKlein(imp: Impostazioni): string { + const vae = sul(imp, 'vae', imp.kleinVae) || sul(imp, 'vae', risoltoSalvato('klein-vae')?.file) + return vae || preferisci(lista(imp, 'vae'), [/^flux2-vae\./i, /klein/i, /flux.?[23]/i, /^ae\./i, /flux/i]) +} + +/** FLUX.2 Klein 4B: il denoiser (FP8 ufficiale), il Qwen3 4B di solo testo e il VAE di FLUX.2 */ +function fileKlein4b(imp: Impostazioni): FileKlein | string { + let modello = sul(imp, 'diffusion_models', risoltoSalvato('klein4b-fp8')?.file) + if (!modello) { + const tutti = [...lista(imp, 'diffusion_models'), ...lista(imp, 'unet')].filter((n) => /klein.*4b|4b.*klein/i.test(n) && !/base/i.test(n) && !/vae|text_enc/i.test(n)) + modello = preferisci(tutti, [/fp8/i, /bf16/i, /./]) + } + const encoder = sul(imp, 'text_encoders', risoltoSalvato('klein4b-te')?.file) || preferisci(lista(imp, 'text_encoders'), [/^qwen_3_4b\.safetensors$/i, /^qwen_?3_4b/i]) + const vae = vaeKlein(imp) + const mancano = [!modello && 'il modello FLUX.2 Klein 4B', !encoder && 'il text encoder (Qwen3 4B)', !vae && 'il VAE'].filter(Boolean) + if (mancano.length) return `FLUX.2 Klein 4B non è pronto: mancano ${mancano.join(', ')}. Scaricali dal selettore del modello (o da Opzioni → Scarica altri modelli).` + return { modello, encoder, vae } +} diff --git a/src/main/motore/lavori.ts b/src/main/motore/lavori.ts index 4f6a405..ce90c46 100644 --- a/src/main/motore/lavori.ts +++ b/src/main/motore/lavori.ts @@ -8,6 +8,7 @@ import type { Famiglia, Impostazioni, Lavoro, Richiesta } from '@shared/tipi' import { USCITA } from '../percorsi' import { impostazioni, profiloMemoria, registraTempo, stimaPasso } from '../impostazioni' import { costruisci, type GrafoPronto, type Ingressi } from './grafi' +import { adattaRichiesta } from './adatta' import { SISTEMA_TRADUCI, daTradurre, pulisciTraduzione, tokenPer } from './traduci' import { accoda, caricaImmagine, chiaviRtx, collegato, connetti, interrompi, nodiDisponibili, scaricaModelli, suAnteprima, suMessaggio, usciteDa, type MessaggioComfy } from './cliente' import { avviaMotore, infoMotore, aggiornaVram } from './processo' @@ -42,7 +43,7 @@ const casuale = (): number => Math.floor(Math.random() * 2 ** 48) /** i LoRA accesi si usano sempre: la loro parola chiave entra nel prompt da sola (se non c'è già), da qualsiasi pagina arrivi la richiesta */ function conParoleLora(q: Richiesta): Richiesta { - if (q.modalita === 'descrivi' || !q.prompt?.trim() || !q.lora?.length || (q.famiglia ?? impostazioni().famiglia) !== 'qwen') return q + if (q.modalita === 'descrivi' || !q.prompt?.trim() || !q.lora?.length) return q let info try { info = elencoLora() @@ -58,7 +59,7 @@ function conParoleLora(q: Richiesta): Richiesta { } export function accodaRichiesta(q0: Richiesta): Lavoro[] { - const q = conParoleLora(q0) + const q = conParoleLora(adattaRichiesta(q0)) const quante = Math.max(1, Math.min(16, q.quante || 1)) const base = q.seed >= 0 ? q.seed : -1 const nuovi: Lavoro[] = [] @@ -276,6 +277,29 @@ async function traduttorePronto(l: Lavoro): Promise { } } +/** le modalità che usano un segmentatore (si scarica da solo, una volta, al primo uso): voce del catalogo e nome in italiano */ +const SEGMENTATORI: Record = { + sfondo: { id: 'sfondo-birefnet', nome: 'BiRefNet' }, + segmenta: { id: 'seg-sam2', nome: 'SAM 2' } +} + +/** il segmentatore è sul disco? se manca lo scarica (BiRefNet ~450 MB, SAM 2 ~185 MB); se non ci riesce il lavoro fallisce con un messaggio chiaro */ +async function segmentatorePronto(l: Lavoro, s: { id: string; nome: string }): Promise { + const stato = (): boolean => !!statoModelli().find((v) => v.id === s.id)?.presente + if (stato()) return + l.fase = `Scarico ${s.nome}` + avvisa(l) + try { + await scaricaVoce(s.id, (r, t) => { + l.fase = `Scarico ${s.nome}${t ? ` · ${Math.round((r / t) * 100)}%` : ''}` + avvisa(l, false) + }) + } catch (e) { + throw new Error(`Non riesco a scaricare ${s.nome}: ${(e as Error).message}`) + } + if (!stato()) throw new Error(`${s.nome} non è sul disco: scaricalo da Opzioni → Scarica altri modelli.`) +} + // le didascalie già scritte (più immagini dallo stesso prompt = un'espansione sola) const potenziati = new Map() @@ -334,11 +358,12 @@ function conEncoderDiTesto(q: Richiesta, imp: Impostazioni): { q: Richiesta; imp } /** il prompt in inglese per il modello: se è in italiano lo traduce il traduttore dedicato (o, se manca, Qwen3-VL già caricato); se no resta com'è */ -async function inInglese(l: Lavoro): Promise { +async function inInglese(l: Lavoro, testo?: string): Promise { const q = l.richiesta - const p = q.prompt.trim() + const originale = testo ?? q.prompt + const p = originale.trim() const imp = impostazioni() - if (!imp.traduci || !CON_PROMPT.has(q.modalita) || !p || !daTradurre(p)) return q.prompt + if (!imp.traduci || !CON_PROMPT.has(q.modalita) || !p || !daTradurre(p)) return originale const gia = tradotti.get(p) if (gia) return gia l.fase = 'Traduco il prompt in inglese' @@ -413,11 +438,12 @@ async function esegui(l: Lavoro): Promise { await avviaMotore() } if (!collegato()) await connetti() + if (l.richiesta.modalita in SEGMENTATORI) await segmentatorePronto(l, SEGMENTATORI[l.richiesta.modalita]) const ing = await prepara(l.richiesta) // FLUX 3 Klein: i suoi file (scelti, scaricati dall'app o messi a mano) e, se serve, il prompt espanso const fam = l.richiesta.famiglia ?? imp.famiglia let impLavoro = imp - const usaKlein = fam !== 'qwen' && l.richiesta.modalita !== 'descrivi' && l.richiesta.modalita !== 'traduci' && l.richiesta.modalita !== 'ingrandisci' + const usaKlein = fam !== 'qwen' && l.richiesta.modalita !== 'descrivi' && l.richiesta.modalita !== 'traduci' && l.richiesta.modalita !== 'ingrandisci' && !(l.richiesta.modalita in SEGMENTATORI) if (usaKlein) { const f = fileKlein(imp, fam) if (typeof f === 'string') throw new Error(f) @@ -436,6 +462,9 @@ async function esegui(l: Lavoro): Promise { } } else inglese = await inInglese(l) if (inglese !== l.richiesta.prompt) l.promptInglese = inglese + // i prompt dei riquadri (Editor box) vanno al modello in inglese come quello della scena + let regioniInglese = l.richiesta.regioni + if (regioniInglese?.length) regioniInglese = await Promise.all(regioniInglese.map(async (r) => ({ ...r, prompt: await inInglese(l, r.prompt) }))) tempiPasso = [] ultimoPasso = 0 // l'ingranditore NVIDIA RTX, se il motore ce l'ha e non l'hai spento in Opzioni @@ -443,14 +472,14 @@ async function esegui(l: Lavoro): Promise { // lavori di testo con Klein: scrivono col suo 4B (vedi conEncoderDiTesto) const lavoroTesto = conEncoderDiTesto(l.richiesta, impLavoro) impLavoro = lavoroTesto.imp - const grafo = costruisci({ ...lavoroTesto.q, prompt: inglese }, impLavoro, ing, l.id.slice(0, 8), profiloMemoria(infoMotore().vramTotale), rtx) + const grafo = costruisci({ ...lavoroTesto.q, prompt: inglese, regioni: regioniInglese }, impLavoro, ing, l.id.slice(0, 8), profiloMemoria(infoMotore().vramTotale), rtx) l.areaAnteprima = grafo.area l.megapixel = grafo.megapixel || undefined l.passiTotali = grafo.passi l.stima = grafo.megapixel ? Math.round(stimaPasso(grafo.megapixel) * grafo.passi + 6) : undefined const testo = l.richiesta.modalita === 'descrivi' // la chiave della VRAM: chi scrive (4B o 8B), chi disegna (la famiglia), o l'ingranditore che non sposta niente - const occupa = testo ? occupaTesto(lavoroTesto.q.encoderKlein) : l.richiesta.modalita === 'ingrandisci' ? 'leggero' : occupaImmagine(fam) + const occupa = testo ? occupaTesto(lavoroTesto.q.encoderKlein) : l.richiesta.modalita === 'ingrandisci' || l.richiesta.modalita in SEGMENTATORI ? 'leggero' : occupaImmagine(fam) const uscita = await lancia(grafo, occupa, l, testo) if (testo) { l.testo = lavoroTesto.q.encoderKlein ? senzaTitolo(uscita?.testo || '') : (uscita?.testo || '').trim() @@ -471,7 +500,7 @@ async function esegui(l: Lavoro): Promise { sampler: q.sampler, scheduler: q.scheduler, lora: q.lora, - modello: usaKlein ? (fam === 'klein-base' ? 'FLUX 3 Klein Base · ' : 'FLUX 3 Klein · ') + impLavoro.kleinModello : imp.modello, + modello: usaKlein ? (fam === 'klein-base' ? 'FLUX 3 Klein Base · ' : fam === 'klein4b' ? 'FLUX.2 Klein 4B · ' : 'FLUX 3 Klein · ') + impLavoro.kleinModello : imp.modello, forza: q.forza, trasparente: q.trasparente, origine: q.origine, diff --git a/src/main/motore/regioni.ts b/src/main/motore/regioni.ts new file mode 100644 index 0000000..12259de --- /dev/null +++ b/src/main/motore/regioni.ts @@ -0,0 +1,78 @@ +// Regional prompting per l'Editor box: dai riquadri disegnati (ognuno col suo prompt) al testo che capisce il modello. +// File puro (niente Electron né import oltre ai tipi): lo prova `node test/regioni.mjs`. +// +// - FLUX 3 Klein (e Klein Base) sa leggere i riquadri: il suo prompt è «didascalia … [righe JSON]», con le righe +// {"id", "bbox": [y_min, x_min, y_max, x_max], "desc"} su una tela 0–1000 (origine in alto a sinistra) e la didascalia che cita ogni +// (formato di black-forest-labs/flux-3-image-klein, src/flux3/bbox.py). Per le modifiche una riga è {"id", "from": "target", "src_bbox": null, "tgt_bbox", "desc"}. +// I riquadri sono una guida morbida: l'elemento può uscire un po' dal rettangolo, e non protegge i pixel fuori (per quello c'è la zona con la maschera). +// - Qwen-Image 2.1 e FLUX.2 Klein 4B non hanno i riquadri: la posizione si dice a parole («at the top left, a small …»). +import type { Box01, Regione } from '../../shared/tipi' + +export interface RegionePronta { + /** id per il modello: obj_1, obj_2… (lettere, cifre e _; il nome scelto dall'interfaccia non conta) */ + id: string + /** [x0, y0, x1, y1] in 0–1, dentro la tela, con almeno 2 % per lato */ + box: Box01 + prompt: string +} + +const MIN = 0.02 +const num = (x: unknown): number => (typeof x === 'number' && Number.isFinite(x) ? x : NaN) + +/** toglie le regioni vuote o senza un riquadro valido, porta i riquadri dentro la tela e dà gli id obj_1… (al massimo 16) */ +export function regioniPronte(regioni: Regione[] | undefined): RegionePronta[] { + const fuori: RegionePronta[] = [] + for (const r of regioni || []) { + const p = (r.prompt || '').trim().replace(/\s+/g, ' ') + if (!p || !Array.isArray(r.box) || r.box.length !== 4) continue + let [x0, y0, x1, y1] = r.box.map(num) + if ([x0, y0, x1, y1].some(Number.isNaN)) continue + if (x1 < x0) [x0, x1] = [x1, x0] + if (y1 < y0) [y0, y1] = [y1, y0] + x0 = Math.max(0, Math.min(1, x0)) + x1 = Math.max(0, Math.min(1, x1)) + y0 = Math.max(0, Math.min(1, y0)) + y1 = Math.max(0, Math.min(1, y1)) + if (x1 - x0 < MIN || y1 - y0 < MIN) continue + fuori.push({ id: `obj_${fuori.length + 1}`, box: [x0, y0, x1, y1], prompt: p }) + if (fuori.length >= 16) break + } + return fuori +} + +/** [x0, y0, x1, y1] 0–1 → [y_min, x_min, y_max, x_max] 0–1000 (come lo vuole Klein) */ +export const bboxKlein = (b: Box01): [number, number, number, number] => [Math.round(b[1] * 1000), Math.round(b[0] * 1000), Math.round(b[3] * 1000), Math.round(b[2] * 1000)] + +const maiuscola = (s: string): string => s.charAt(0).toUpperCase() + s.slice(1) +const senzaPunto = (s: string): string => s.trim().replace(/[.\s]+$/, '') + +/** il prompt di Klein con i riquadri: didascalia che cita ogni , poi le righe JSON. `modifica`: i riquadri dicono dove va quello che si aggiunge/cambia. */ +export function promptKleinConBox(scena: string, regioni: RegionePronta[], modifica = false): string { + if (!regioni.length) return scena + const righe = regioni.map((r) => + modifica ? { id: r.id, from: 'target', src_bbox: null, tgt_bbox: bboxKlein(r.box), desc: maiuscola(r.prompt) } : { id: r.id, bbox: bboxKlein(r.box), desc: maiuscola(r.prompt) } + ) + const testi = regioni.map((r) => `${maiuscola(senzaPunto(r.prompt))} <${r.id}>.`) + const didascalia = [scena.trim() ? maiuscola(senzaPunto(scena)) + '.' : '', ...testi].filter(Boolean).join(' ') + return `${didascalia} ${JSON.stringify(righe)}` +} + +/** dov'è un riquadro, a parole (per i modelli che non leggono i riquadri) */ +export function dove(b: Box01): string { + const cx = (b[0] + b[2]) / 2 + const cy = (b[1] + b[3]) / 2 + const o = cx < 0.34 ? 'left' : cx > 0.66 ? 'right' : '' + const v = cy < 0.34 ? 'top' : cy > 0.66 ? 'bottom' : '' + const area = (b[2] - b[0]) * (b[3] - b[1]) + const grande = area > 0.4 ? 'filling most of the frame, ' : area > 0.16 ? 'large, ' : area < 0.04 ? 'small, ' : '' + const posto = o || v ? `at the ${[v, o].filter(Boolean).join(' ')} of the frame` : 'in the center of the frame' + return grande + posto +} + +/** il prompt per Qwen-Image 2.1 e FLUX.2 Klein 4B: la posizione di ogni elemento a parole */ +export function promptConPosizioni(scena: string, regioni: RegionePronta[], modifica = false): string { + if (!regioni.length) return scena + const testi = regioni.map((r) => (modifica ? `${maiuscola(senzaPunto(r.prompt))} — ${dove(r.box)}.` : `${maiuscola(dove(r.box))}: ${senzaPunto(r.prompt)}.`)) + const apertura = scena.trim() ? maiuscola(senzaPunto(scena)) + '. ' : '' + return `${apertura}${modifica ? 'Apply these changes, each one in its place:' : 'Compose the image with these elements, each in its place:'} ${testi.join(' ')}` +} diff --git a/src/main/motore/spazio.ts b/src/main/motore/spazio.ts index dc76b8d..ade5e13 100644 --- a/src/main/motore/spazio.ts +++ b/src/main/motore/spazio.ts @@ -7,10 +7,10 @@ import type { Famiglia } from '../../shared/tipi' * - `testo:4b` / `testo:8b`: l'LLM che scrive (il Qwen3-VL 4B di Klein o l'8B di Qwen); * - `cpu`: il traduttore dedicato, che sta sul processore; * - `leggero`: l'ingranditore (pochi MB): non chiede spazio e non cambia quello che c'è. */ -export type Occupa = 'immagine:qwen' | 'immagine:klein' | 'immagine:klein-base' | 'testo:4b' | 'testo:8b' | 'cpu' | 'leggero' +export type Occupa = 'immagine:qwen' | 'immagine:klein' | 'immagine:klein-base' | 'immagine:klein4b' | 'testo:4b' | 'testo:8b' | 'cpu' | 'leggero' /** la chiave di un lavoro che disegna */ -export const occupaImmagine = (fam: Famiglia): Occupa => (fam === 'klein' ? 'immagine:klein' : fam === 'klein-base' ? 'immagine:klein-base' : 'immagine:qwen') +export const occupaImmagine = (fam: Famiglia): Occupa => (fam === 'klein' ? 'immagine:klein' : fam === 'klein4b' ? 'immagine:klein4b' : fam === 'klein-base' ? 'immagine:klein-base' : 'immagine:qwen') /** la chiave di un lavoro che scrive: col 4B di Klein (`encoderKlein`) o con l'8B di Qwen */ export const occupaTesto = (encoderKlein?: boolean): Occupa => (encoderKlein ? 'testo:4b' : 'testo:8b') @@ -40,6 +40,8 @@ export function faseScarico(prima: Occupa | null): string { return 'Scarico FLUX 3 Klein' case 'immagine:klein-base': return 'Scarico FLUX 3 Klein Base' + case 'immagine:klein4b': + return 'Scarico FLUX.2 Klein 4B' case 'testo:4b': case 'testo:8b': return 'Scarico l\'LLM' diff --git a/src/preload/index.ts b/src/preload/index.ts index d47dee5..1c2432a 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -28,7 +28,7 @@ const api = { impostazioni: { leggi: chiama('impostazioni:leggi'), salva: chiama('impostazioni:salva') }, sistema: { controllo: chiama('sistema:controllo') }, setup: { cerca: chiama('setup:cerca'), stato: chiama('setup:stato'), avvia: chiama('setup:avvia'), annulla: chiama('setup:annulla') }, - modelli: { stato: chiama('modelli:stato'), file: chiama('modelli:file'), scarica: chiama('modelli:scarica'), annulla: chiama('download:annulla'), inutilizzati: chiama('modelli:inutilizzati'), elimina: chiama('modelli:elimina') }, + modelli: { stato: chiama('modelli:stato'), file: chiama('modelli:file'), scarica: chiama('modelli:scarica'), annulla: chiama('download:annulla'), segmentatori: chiama('modelli:segmentatori'), inutilizzati: chiama('modelli:inutilizzati'), elimina: chiama('modelli:elimina') }, motore: { info: chiama('motore:info'), avvia: chiama('motore:avvia'), riavvia: chiama('motore:riavvia'), ferma: chiama('motore:ferma'), log: chiama('motore:log'), libera: chiama('motore:libera'), apriComfy: chiama('motore:apriComfy') @@ -40,7 +40,7 @@ const api = { mostra: chiama('galleria:mostra'), apriCartella: chiama('galleria:apriCartella'), copia: chiama('galleria:copia'), esporta: chiama('galleria:esporta'), importa: chiama('galleria:importa'), meta: chiama('galleria:meta'), importaFoto: chiama('galleria:importaFoto') }, - lora: { elenco: chiama('lora:elenco'), aggiorna: chiama('lora:aggiorna'), importa: chiama('lora:importa'), elimina: chiama('lora:elimina'), scarica: chiama('lora:scarica'), apriCartella: chiama('lora:apriCartella') }, + lora: { elenco: chiama('lora:elenco'), compatibili: chiama('lora:compatibili'), aggiorna: chiama('lora:aggiorna'), importa: chiama('lora:importa'), elimina: chiama('lora:elimina'), scarica: chiama('lora:scarica'), apriCartella: chiama('lora:apriCartella') }, file: { scegliImmagini: chiama('file:scegliImmagini'), scegli: chiama('file:scegli'), scegliCartella: chiama('file:scegliCartella'), salvaTemp: chiama('file:salvaTemp'), info: chiama('file:info'), portaDentro: chiama('file:portaDentro'), inPng: chiama('file:inPng'), diff --git a/src/renderer/src/api.ts b/src/renderer/src/api.ts index ae61b6e..e49f8ce 100644 --- a/src/renderer/src/api.ts +++ b/src/renderer/src/api.ts @@ -1,6 +1,6 @@ // L'API del processo principale, con i tipi giusti. import type { - ControlloSistema, Download, FiltroGalleria, Impostazioni, InfoLora, InfoMotore, Lavoro, Opera, PassoSetup, Richiesta, Risorse, StatoAggiornamento, StatoVoce + ControlloSistema, Download, Famiglia, FiltroGalleria, Impostazioni, InfoLora, InfoMotore, Lavoro, Opera, PassoSetup, Richiesta, Risorse, StatoAggiornamento, StatoVoce } from '@shared/tipi' type Via = () => void @@ -36,6 +36,9 @@ export interface FileModelli { klein: string[] kleinTe: string[] kleinVae: string[] + /** i file di FLUX.2 Klein 4B (modello, text encoder, VAE condiviso) */ + klein4b: string[] + klein4bTe: string[] } export interface ModelloTrovato { percorso: string @@ -63,6 +66,8 @@ export interface Api { modelli: { stato(): Promise; file(): Promise; scarica(id: string): Promise; annulla(id: string): Promise inutilizzati(): Promise; elimina(percorso: string): Promise + /** BiRefNet (togli sfondo) e SAM 2 (estrai fedele): `presente` false = si scarica da solo al primo uso */ + segmentatori(): Promise } motore: { info(): Promise; avvia(): Promise; riavvia(): Promise; ferma(): Promise @@ -77,7 +82,9 @@ export interface Api { importaFoto(png: string, nome: string): Promise } lora: { - elenco(): Promise; aggiorna(file: string, m: Partial): Promise; importa(p: string[]): Promise + elenco(): Promise; aggiorna(file: string, m: Partial): Promise; + /** i LoRA buoni per quel modello (compatibili e «da provare»; mai quelli fatti per un altro) */ + compatibili(famiglia: Famiglia): Promise; importa(p: string[]): Promise elimina(file: string): Promise; scarica(link: string): Promise; apriCartella(): Promise } file: { diff --git a/src/renderer/src/componenti/ScegliModello.tsx b/src/renderer/src/componenti/ScegliModello.tsx index d20abb3..1bde24e 100644 --- a/src/renderer/src/componenti/ScegliModello.tsx +++ b/src/renderer/src/componenti/ScegliModello.tsx @@ -9,16 +9,16 @@ import { ricaricaImp, usaStato } from '../stato' import { I } from './Icone' /** le voci del catalogo da avere per usare un modello (il modello INT8 sta nelle schede da 8 GB) */ -const SERVONO: Record, string[]> = { klein: ['klein-int8', 'klein-te', 'klein-vae'], 'klein-base': ['klein-base-int8', 'klein-te', 'klein-vae'] } -const PESO: Record, string> = { klein: 'circa 10 GB', 'klein-base': 'circa 10 GB (4,7 se hai già Klein)' } -const ETICHETTA: Record = { 'qwen-base': 'Qwen base', klein: 'Klein veloce', 'klein-base': 'Klein Base' } +const SERVONO: Record, string[]> = { klein: ['klein-int8', 'klein-te', 'klein-vae'], 'klein-base': ['klein-base-int8', 'klein-te', 'klein-vae'], klein4b: ['klein4b-fp8', 'klein4b-te', 'klein-vae'] } +const PESO: Record, string> = { klein: 'circa 10 GB', 'klein-base': 'circa 10 GB (4,7 se hai già Klein)', klein4b: 'circa 12,5 GB' } +const ETICHETTA: Record = { 'qwen-base': 'Qwen base', klein: 'Klein veloce', 'klein-base': 'Klein Base', klein4b: 'Klein 4B' } export function SelettoreModello(): JSX.Element { const imp = usaStato((s) => s.imp) const avvisa = usaStato((s) => s.avvisa) const inCorso = usaStato((s) => s.lavori.some((l) => l.stato === 'in coda' || l.stato === 'in corso')) const scelto = modelloDaImpostazioni(imp) - const [pronti, setPronti] = useState>({ 'qwen-base': true, klein: false, 'klein-base': false }) + const [pronti, setPronti] = useState>({ 'qwen-base': true, klein: false, 'klein-base': false, klein4b: false }) const [aperto, setAperto] = useState(null) const [scarico, setScarico] = useState<{ id: IdModello; passo: string; perc: number } | null>(null) const radice = useRef(null) @@ -30,7 +30,8 @@ export function SelettoreModello(): JSX.Element { setPronti({ 'qwen-base': true, klein: comuni && !!f.klein?.some((n) => !/base/i.test(n)), - 'klein-base': comuni && !!f.klein?.some((n) => /base/i.test(n)) + 'klein-base': comuni && !!f.klein?.some((n) => /base/i.test(n)), + klein4b: (f.klein4b?.length ?? 0) > 0 && (f.klein4bTe?.length ?? 0) > 0 && (f.kleinVae?.length ?? 0) > 0 }) } useEffect(() => { @@ -68,7 +69,7 @@ export function SelettoreModello(): JSX.Element { const scarica = async (id: Exclude): Promise => { try { const stato = await api.modelli.stato() - const nomi: Record = { 'klein-int8': 'il modello Klein', 'klein-base-int8': 'il modello Klein Base', 'klein-te': 'il text encoder', 'klein-vae': 'il VAE' } + const nomi: Record = { 'klein-int8': 'il modello Klein', 'klein-base-int8': 'il modello Klein Base', 'klein4b-fp8': 'il modello Klein 4B', 'klein4b-te': 'il text encoder (Qwen3 4B)', 'klein-te': 'il text encoder', 'klein-vae': 'il VAE' } for (const voce of SERVONO[id]) { if (stato.find((v) => v.id === voce)?.presente) continue setScarico({ id, passo: nomi[voce], perc: 0 }) diff --git a/src/renderer/src/pagine/Impostazioni.tsx b/src/renderer/src/pagine/Impostazioni.tsx index b6d8123..ec6cc19 100644 --- a/src/renderer/src/pagine/Impostazioni.tsx +++ b/src/renderer/src/pagine/Impostazioni.tsx @@ -29,7 +29,7 @@ export function Impostazioni(): JSX.Element { function OpzioniGenerali(): JSX.Element { const { imp, motore, avvisa, agg } = usaStato() - const [file, setFile] = useState({ diffusione: [], encoder: [], vae: [], upscaler: [], klein: [], kleinTe: [], kleinVae: [] }) + const [file, setFile] = useState({ diffusione: [], encoder: [], vae: [], upscaler: [], klein: [], kleinTe: [], kleinVae: [], klein4b: [], klein4bTe: [] }) const [catalogo, setCatalogo] = useState([]) const [download, setDownload] = useState>({}) const [log, setLog] = useState([]) diff --git a/src/shared/loraBase.ts b/src/shared/loraBase.ts new file mode 100644 index 0000000..13b7543 --- /dev/null +++ b/src/shared/loraBase.ts @@ -0,0 +1,48 @@ +// Da quali nomi dei tensori di un LoRA si capisce per quale modello è fatto. File puro (niente Electron): lo prova `node test/lora.mjs`. +// Qwen-Image 2.1 transformer_blocks.N.attn.to_q / img_mlp (flusso singolo) +// Qwen-Image 1.x txt_mlp, add_q_proj… (due flussi) +// FLUX 3 Klein single_blocks.N.(qkv|q_proj|attn_out|mlp_in|mlp_out), content_mode_blocks, txt_mode_blocks (ComfyUI li mappa: comfy/lora.py, ramo Flux3) +// FLUX.2 Klein 4B double_blocks.N.(img_attn|txt_attn|img_mlp|txt_mlp) + single_blocks.N.(linear1|linear2), 5 doppi blocchi e dimensione 3072 +// FLUX.1 / FLUX.2 grandi / SD / SDXL: non compatibili con nessuno dei modelli dell'app +import type { Famiglia, InfoLora } from './tipi' + +export interface AnalisiLora { + compatibile: InfoLora['compatibile'] + base?: string + /** per quali famiglie dell'app è buono (vuoto = nessuna / non si capisce) */ + famiglie: Famiglia[] +} + +const SEP = '[._]' +const blocco = (nome: string, resto: string): RegExp => new RegExp(`${nome}${SEP}\\d+${SEP}(?:${resto})`) + +/** `chiavi`: i nomi dei tensori; `forme`: (facoltativo) le forme, per distinguere FLUX.2 4B da 9B */ +export function analizzaChiavi(chiavi: string[], forme?: Record): AnalisiLora { + const tutte = chiavi.join('\n') + if (!chiavi.length) return { compatibile: '?', famiglie: [] } + if (/txt_mlp|add_q_proj|txt_mod|img_mod\./.test(tutte) && !/double_blocks/.test(tutte)) return { compatibile: 'no', base: 'Qwen-Image 1.x', famiglie: [] } + // FLUX 3 Klein: i blocchi «modali» sono solo suoi; `single_blocks` con qkv/attn_out/mlp_in/mlp_out pure (FLUX.1 e FLUX.2 usano linear1/linear2) + if (/(content_mode_blocks|txt_mode_blocks)/.test(tutte) || blocco('single_blocks', 'qkv|q_proj|k_proj|v_proj|attn_out|mlp_in|mlp_out').test(tutte)) { + return { compatibile: 'si', base: 'FLUX 3 Klein', famiglie: ['klein', 'klein-base'] } + } + if (/(?:double|single)_blocks/.test(tutte)) { + // FLUX.2 Klein 4B: 5 blocchi doppi e dimensione 3072; il resto (FLUX.1, FLUX.2 9B/dev) non è dell'app + const doppi = chiavi.map((k) => /double_blocks[._](\d+)/.exec(k)?.[1]).filter((x): x is string => x !== undefined).map(Number) + const maxDoppio = doppi.length ? Math.max(...doppi) : -1 + const chiaveIn = chiavi.find((k) => /double_blocks[._]0[._]img_attn[._]qkv.*(lora_A|lora_down)/.test(k)) + const dimIn = chiaveIn ? forme?.[chiaveIn]?.[1] : undefined + const flux2 = /linear1/.test(tutte) && /img_mlp|txt_mlp/.test(tutte) + if (maxDoppio >= 0 && maxDoppio <= 4 && (dimIn === undefined || dimIn === 3072) && flux2) return { compatibile: 'si', base: 'FLUX.2 Klein 4B', famiglie: ['klein4b'] } + return { compatibile: 'no', base: maxDoppio <= 4 ? 'FLUX.2' : 'FLUX.1 / FLUX.2 grande', famiglie: [] } + } + if (/lora_unet_|down_blocks|input_blocks|lora_te/.test(tutte)) return { compatibile: 'no', base: 'Stable Diffusion / SDXL', famiglie: [] } + if (/transformer_blocks\.\d+\.(attn\.to_q|img_mlp)/.test(tutte)) return { compatibile: 'si', base: 'Qwen-Image 2.1', famiglie: ['qwen'] } + return { compatibile: '?', famiglie: [] } +} + +/** un LoRA va bene con questo modello? 'si' (provato dai nomi), 'no', o '?' (non si capisce: si lascia provare) */ +export function vaBeneCon(l: { compatibile: InfoLora['compatibile']; famiglie?: Famiglia[] }, fam: Famiglia): 'si' | 'no' | '?' { + if (l.compatibile === '?') return '?' + if (!l.famiglie || !l.famiglie.length) return l.compatibile === 'si' ? 'si' : 'no' + return l.famiglie.includes(fam) ? 'si' : 'no' +} diff --git a/src/shared/modelli.ts b/src/shared/modelli.ts index ce62a7f..ce36b84 100644 --- a/src/shared/modelli.ts +++ b/src/shared/modelli.ts @@ -9,7 +9,7 @@ // I grafi di Qwen non cambiano: qwen-base manda la stessa richiesta di prima con la qualità «Alta». import type { Famiglia } from './tipi' -export type IdModello = 'qwen-base' | 'klein' | 'klein-base' +export type IdModello = 'qwen-base' | 'klein' | 'klein-base' | 'klein4b' export interface InfoModello { id: IdModello @@ -30,12 +30,13 @@ export interface InfoModello { export const MODELLI: Record = { 'qwen-base': { id: 'qwen-base', famiglia: 'qwen', turbo: false, passi: 40, cfg: null, nome: 'Qwen-Image 2.1', corto: 'Qwen', sotto: '40 passi · LoRA', lora: true, veloce: false }, - klein: { id: 'klein', famiglia: 'klein', turbo: false, passi: 8, cfg: 1, nome: 'FLUX 3 Klein', corto: 'Klein veloce', sotto: '8 passi · senza CFG · il più veloce', lora: false, veloce: true }, - 'klein-base': { id: 'klein-base', famiglia: 'klein-base', turbo: false, passi: 50, cfg: 3.5, nome: 'FLUX 3 Klein Base', corto: 'Klein Base', sotto: '50 passi · CFG 3,5 · la qualità piena', lora: false, veloce: false } + klein: { id: 'klein', famiglia: 'klein', turbo: false, passi: 8, cfg: 1, nome: 'FLUX 3 Klein', corto: 'Klein veloce', sotto: '8 passi · senza CFG · il più veloce', lora: true, veloce: true }, + 'klein-base': { id: 'klein-base', famiglia: 'klein-base', turbo: false, passi: 50, cfg: 3.5, nome: 'FLUX 3 Klein Base', corto: 'Klein Base', sotto: '50 passi · CFG 3,5 · la qualità piena', lora: true, veloce: false }, + klein4b: { id: 'klein4b', famiglia: 'klein4b', turbo: false, passi: 4, cfg: 1, nome: 'FLUX.2 Klein 4B', corto: 'Klein 4B', sotto: '4 passi · Apache-2.0 · uso commerciale libero', lora: true, veloce: true } } /** l'ordine delle pillole nell'interfaccia */ -export const ORDINE_MODELLI: IdModello[] = ['qwen-base', 'klein', 'klein-base'] +export const ORDINE_MODELLI: IdModello[] = ['qwen-base', 'klein', 'klein-base', 'klein4b'] export const MODELLO_PREDEFINITO: IdModello = 'qwen-base' export const eModello = (x: unknown): x is IdModello => typeof x === 'string' && Object.prototype.hasOwnProperty.call(MODELLI, x) @@ -44,6 +45,7 @@ export const eModello = (x: unknown): x is IdModello => typeof x === 'string' && export function modelloDa(famiglia?: unknown, _turbo?: boolean): IdModello { if (famiglia === 'klein') return 'klein' if (famiglia === 'klein-base') return 'klein-base' + if (famiglia === 'klein4b') return 'klein4b' return 'qwen-base' } @@ -69,6 +71,7 @@ export function allineaScelta(modifiche: { modelloScelto?: unknown; famiglia?: u * «FLUX 3 Klein Base · file», «FLUX 3 Klein · file» o il file del denoiser di Qwen (Turbo compreso, che ora vale Qwen base) */ export function modelloDaOpera(o: { modello?: string | null } | null | undefined): IdModello { const m = (o?.modello || '').trim() + if (/^FLUX\.2 Klein 4B\b/i.test(m)) return 'klein4b' if (/^FLUX 3 Klein Base\b/i.test(m)) return 'klein-base' if (/^FLUX 3 Klein\b/i.test(m)) return 'klein' return 'qwen-base' diff --git a/src/shared/tipi.ts b/src/shared/tipi.ts index 6a39c53..5b6925c 100644 --- a/src/shared/tipi.ts +++ b/src/shared/tipi.ts @@ -14,12 +14,53 @@ export type Modalita = | 'descrivi' // il modello di testo scrive (migliora prompt, descrive una foto, traduce) | 'traduci' // il prompt dall'italiano all'inglese con il traduttore dedicato (sul processore, senza Qwen) | 'importa' // (solo in galleria) una foto portata dentro da fuori + | 'sfondo' // toglie lo sfondo: PNG con trasparenza (BiRefNet), i pixel del soggetto restano quelli originali + | 'segmenta' // estrae da una foto la cosa indicata (riquadro o punti): SAM2 → PNG con trasparenza, pixel originali export interface LoraAttiva { file: string + /** −2…2 (0 = ignorato); l'ordine della lista è l'ordine in cui si applicano */ forza: number } +/** Un riquadro dell'Editor box: coordinate 0–1 sulla tela/foto (origine in alto a sinistra), [x0, y0, x1, y1]. */ +export type Box01 = [number, number, number, number] + +/** Una regione dell'Editor box, col suo prompt (regional prompting) */ +export interface Regione { + id: string + box: Box01 + prompt: string +} + +/** Un riquadro dell'Editor box della GUI v2: dove sta nella tela finale (frazioni 0–1) e cosa ci deve essere. + * Il motore lo trasforma in `Regione` (`adattaRichiesta`): la conversione al formato di Klein 3 o alle parole di Qwen la fa solo lui. */ +export interface BoxRegione { + id: string + /** x0, y0, x1, y1 normalizzati 0–1 sulla tela finale */ + rect: [number, number, number, number] + tipo: 'persona' | 'oggetto' | 'testo' | 'sfondo' + prompt: string + /** foto di riferimento di questo box (→ ): il motore non la usa ancora */ + rif?: string + /** guida = composizione suggerita · decisa = il modello resta lì dentro · blocca = maschera e compositing (per ora vale come decisa) */ + rigidita: 'guida' | 'decisa' | 'blocca' + ordine: number + /** l'occhio spento nell'editor: il riquadro resta ma non parte */ + nascosta?: boolean +} + +/** Lo sfondo di una creazione: la scena, trasparente (PNG con alfa) o una tinta unita (#rrggbb). */ +export type SfondoRichiesto = 'scena' | 'trasparente' | { tinta: string } + +/** Estrai fedele: cosa ritagliare con SAM2 (tutto in coordinate 0–1 sulla foto) */ +export interface OpzioniSegmenta { + box?: Box01 + punti?: { x: number; y: number; positivo: boolean }[] + /** true = il PNG è ritagliato attorno alla cosa (+ un margine); false = ha la misura della foto, con il resto trasparente */ + ritaglia?: boolean +} + export interface Bordi { sinistra: number sopra: number @@ -34,8 +75,8 @@ export interface Riquadro { h: number } -/** il modello che disegna: Qwen-Image 2.1, o FLUX 3 Klein (veloce, distillato) o FLUX 3 Klein Base (50 passi con CFG) */ -export type Famiglia = 'qwen' | 'klein' | 'klein-base' +/** il modello che disegna: Qwen-Image 2.1, FLUX 3 Klein (veloce, distillato), FLUX 3 Klein Base (50 passi con CFG) o FLUX.2 Klein 4B (Apache-2.0, 4 passi) */ +export type Famiglia = 'qwen' | 'klein' | 'klein-base' | 'klein4b' /** quello che distingue i modelli: passi per qualità, CFG e cosa sanno fare. Nomi e passi «alti» vengono da `modelli.ts` * (la mappa unica); «turbo» resta solo per le qualità salvate prima della 1.1 e vale come «alta» (`senzaTurbo`). */ @@ -46,7 +87,8 @@ const famiglia = (id: IdModello, bozza: number, massima: number): { nome: string export const FAMIGLIE: Record> = { qwen: famiglia('qwen-base', 20, 50), klein: famiglia('klein', 8, 16), - 'klein-base': famiglia('klein-base', 25, 64) + 'klein-base': famiglia('klein-base', 25, 64), + klein4b: famiglia('klein4b', 4, 8) } /** Quello che l'interfaccia chiede: il principale lo trasforma nel grafo di ComfyUI. */ @@ -94,6 +136,15 @@ export interface Richiesta { /** rifinisci/ingrandisci: di quanto ingrandire */ fattore?: number lora: LoraAttiva[] + /** Editor box (regional prompting): più riquadri, ognuno col suo prompt; `prompt` resta la descrizione di tutta la scena. Per 'crea' e 'modifica'. */ + regioni?: Regione[] + /** GUI v2, Editor box: i riquadri con il loro testo (→ `regioni`) e la scena generale (→ `prompt`, se ci sono riquadri validi) */ + box?: BoxRegione[] + promptScena?: string + /** GUI v2: lo sfondo richiesto ('trasparente' = come `trasparente`; { tinta } = una riga in più nel prompt) */ + sfondo?: SfondoRichiesto + /** 'segmenta': cosa estrarre */ + segmenta?: OpzioniSegmenta /** Turbo: LoRA di distillazione a 8 passi (CFG 1), circa 5 volte più veloce */ turbo?: boolean trasparente?: boolean @@ -190,6 +241,8 @@ export interface InfoLora { /** 'si' = Qwen-Image 2.1, 'no' = un altro modello, '?' = non si capisce */ compatibile: 'si' | 'no' | '?' base?: string + /** per quali modelli è fatto (da `analizza`: dai nomi dei tensori); vuoto = non si capisce */ + famiglie?: Famiglia[] aggiunta: number } @@ -255,6 +308,12 @@ export interface Impostazioni { testoCon4B: boolean /** secondi senza lavori dopo i quali i modelli escono dalla VRAM (0 = mai) */ scaricaDopo: number + /** interruttore «Modelli uncensored» (Opzioni): spento di serie. Con lui spento le voci `uncensored` del catalogo non si scaricano */ + uncensored: boolean + /** ha letto e accettato l'avviso sui modelli uncensored (serve anche questo per scaricarli) */ + uncensoredAccettato: boolean + /** i lavori di testo con Klein (migliora prompt, descrivi foto, Idee…) scrivono col Qwen3-VL 4B uncensored (voce `llm-huihui-4b`), se è scaricato; serve `uncensored` acceso */ + testoUncensored: boolean } /** RAM, VRAM e GPU per la barra in alto (byte; uso e temperatura in % e °C) */ @@ -283,7 +342,7 @@ export interface InfoMotore { /** Un file di modello che si può scaricare. */ export interface VoceCatalogo { id: string - tipo: 'diffusione' | 'encoder' | 'vae' | 'lora' | 'upscaler' | 'anteprima' | 'traduttore' | 'klein' | 'klein-te' | 'klein-vae' + tipo: 'diffusione' | 'encoder' | 'vae' | 'lora' | 'upscaler' | 'anteprima' | 'traduttore' | 'klein' | 'klein-te' | 'klein-vae' | 'sfondo' | 'segmentatore' | 'llm' nome: string descrizione: string file: string @@ -315,6 +374,17 @@ export interface VoceCatalogo { famiglia?: string /** non so ancora il nome vero del file: lo trovo al primo download */ daCercare?: boolean + /** il componente è «uncensored» (senza il rifiuto addestrato nei pesi): si scarica solo con l'interruttore acceso e l'avviso accettato */ + uncensored?: boolean + /** sigla della licenza dei pesi (es. 'Apache-2.0', 'Qwen Research', 'FLUX Kommunity', 'MIT') e dove si legge */ + licenza?: string + licenzaUrl?: string + /** a cosa serve nella pipeline: la UI lo mostra accanto al nome */ + ruolo?: 'denoiser' | 'encoder' | 'llm' | 'vae' | 'lora' | 'segmentatore' | 'sfondo' + /** una riga di avvertenza da mostrare prima del download (limiti d'uso, compatibilità) */ + avviso?: string + /** il repo e la revisione fissata (per riprodurre il download); senza revisione si usa `main` */ + revisione?: string } export interface StatoVoce extends VoceCatalogo { diff --git a/test/adatta.mjs b/test/adatta.mjs new file mode 100644 index 0000000..ae791fd --- /dev/null +++ b/test/adatta.mjs @@ -0,0 +1,40 @@ +// Test dell'adattatore GUI v2 → motore (src/main/motore/adatta.ts): box → regioni, sfondo → trasparente / tinta. +// +// Esecuzione: node test/adatta.mjs +import { fileURLToPath, pathToFileURL } from 'node:url' +import { dirname, join } from 'node:path' + +const QUI = dirname(fileURLToPath(import.meta.url)) +const A = await import(pathToFileURL(join(QUI, '..', 'src', 'main', 'motore', 'adatta.ts')).href) + +let falliti = 0 +let fatti = 0 +function uguale(nome, avuto, atteso) { + fatti++ + const a = JSON.stringify(avuto) + const b = JSON.stringify(atteso) + if (a === b) return + falliti++ + console.log(`NO ${nome}\n avuto: ${a}\n atteso: ${b}`) +} +const q = { modalita: 'crea', prompt: 'old flat prompt', lora: [] } +const bx = (id, ordine, o = {}) => ({ id, rect: [0, 0, 0.5, 0.5], tipo: 'oggetto', prompt: 'p' + id, rigidita: 'decisa', ordine, ...o }) + +uguale('senza box né sfondo: invariata', A.adattaRichiesta(q), q) +const r = A.adattaRichiesta({ ...q, promptScena: 'a harbour at dusk', box: [bx('b', 2), bx('a', 1), bx('n', 0, { nascosta: true }), bx('v', 3, { prompt: ' ' })] }) +uguale('box → regioni nell\'ordine, senza nascoste né vuote', r.regioni.map((x) => x.id), ['a', 'b']) +uguale('box → riquadro e prompt', r.regioni[0], { id: 'a', box: [0, 0, 0.5, 0.5], prompt: 'pa' }) +uguale('la scena diventa il prompt', r.prompt, 'a harbour at dusk') +uguale('box senza promptScena: il prompt resta', A.adattaRichiesta({ ...q, box: [bx('a', 0)] }).prompt, 'old flat prompt') +uguale('tutti nascosti: prompt e regioni non cambiano', A.adattaRichiesta({ ...q, promptScena: 'x', box: [bx('a', 0, { nascosta: true })] }), { ...q, promptScena: 'x', box: [bx('a', 0, { nascosta: true })] }) +uguale('regioni già presenti si tengono', A.adattaRichiesta({ ...q, regioni: [{ id: 'z', box: [0, 0, 1, 1], prompt: 'zz' }], box: [bx('a', 0)] }).regioni.map((x) => x.id), ['z', 'a']) + +uguale('sfondo trasparente', A.adattaRichiesta({ ...q, sfondo: 'trasparente' }).trasparente, true) +uguale('sfondo scena: niente', A.adattaRichiesta({ ...q, sfondo: 'scena' }), { ...q, sfondo: 'scena' }) +uguale('tinta unita nel prompt', A.adattaRichiesta({ ...q, sfondo: { tinta: '#ff8800' } }).prompt, 'old flat prompt, on a plain solid #ff8800 background') +uguale('tinta senza #', A.adattaRichiesta({ ...q, sfondo: { tinta: 'ABC' } }).prompt, 'old flat prompt, on a plain solid #ABC background') +uguale('tinta non valida: ignorata', A.adattaRichiesta({ ...q, sfondo: { tinta: 'rosso' } }).prompt, 'old flat prompt') +uguale('tinta con punto finale nel prompt', A.adattaRichiesta({ ...q, prompt: 'a cat.', sfondo: { tinta: '#000' } }).prompt, 'a cat, on a plain solid #000 background') + +console.log(falliti ? `${falliti} FALLITI su ${fatti}` : `TUTTI OK (${fatti} controlli)`) +process.exit(falliti ? 1 : 0) diff --git a/test/campioni-hf/AtomicChat__Qwen-Image-2.1-Turbo-GGUF.json b/test/campioni-hf/AtomicChat__Qwen-Image-2.1-Turbo-GGUF.json new file mode 100644 index 0000000..8d0cb4f --- /dev/null +++ b/test/campioni-hf/AtomicChat__Qwen-Image-2.1-Turbo-GGUF.json @@ -0,0 +1,88 @@ +{ + "repo": "AtomicChat/Qwen-Image-2.1-Turbo-GGUF", + "revisione": "main", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 2447 + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-AD-Q2_K.gguf", + "size": 2553333472, + "lfs": { + "size": 2553333472 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-AD-Q3_K.gguf", + "size": 3599812320, + "lfs": { + "size": 3599812320 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-AD-Q4_K.gguf", + "size": 4201694944, + "lfs": { + "size": 4201694944 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-AD-Q5_K.gguf", + "size": 5399168736, + "lfs": { + "size": 5399168736 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-AD-Q6_K.gguf", + "size": 6711985888, + "lfs": { + "size": 6711985888 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-BF16.gguf", + "size": 14230275808, + "lfs": { + "size": 14230275808 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-Q8_0.gguf", + "size": 7591554784, + "lfs": { + "size": 7591554784 + } + }, + { + "type": "file", + "path": "README.md", + "size": 11653 + }, + { + "type": "file", + "path": "quality-vs-size.png", + "size": 108564, + "lfs": { + "size": 108564 + } + }, + { + "type": "file", + "path": "samples-grid.png", + "size": 3587494, + "lfs": { + "size": 3587494 + } + } + ] +} \ No newline at end of file diff --git a/test/campioni-hf/AtomicChat__Qwen-Image-2.1-Turbo-Uncensored-GGUF.json b/test/campioni-hf/AtomicChat__Qwen-Image-2.1-Turbo-Uncensored-GGUF.json new file mode 100644 index 0000000..8b3cc4e --- /dev/null +++ b/test/campioni-hf/AtomicChat__Qwen-Image-2.1-Turbo-Uncensored-GGUF.json @@ -0,0 +1,88 @@ +{ + "repo": "AtomicChat/Qwen-Image-2.1-Turbo-Uncensored-GGUF", + "revisione": "main", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 3179 + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-Uncensored-AD-Q2_K.gguf", + "size": 3336289216, + "lfs": { + "size": 3336289216 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-Uncensored-AD-Q3_K.gguf", + "size": 4184140736, + "lfs": { + "size": 4184140736 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-Uncensored-AD-Q4_K.gguf", + "size": 5287517120, + "lfs": { + "size": 5287517120 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-Uncensored-AD-Q5_K.gguf", + "size": 6281292736, + "lfs": { + "size": 6281292736 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-Uncensored-AD-Q6_K.gguf", + "size": 7490550720, + "lfs": { + "size": 7490550720 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-Uncensored-BF16.gguf", + "size": 16388045536, + "lfs": { + "size": 16388045536 + } + }, + { + "type": "file", + "path": "Qwen-Image-2.1-Turbo-Uncensored-Q8_0.gguf", + "size": 8709520096, + "lfs": { + "size": 8709520096 + } + }, + { + "type": "file", + "path": "README.md", + "size": 12027 + }, + { + "type": "file", + "path": "mmproj-Qwen-Image-2.1-Turbo-Uncensored-BF16.gguf", + "size": 1162568928, + "lfs": { + "size": 1162568928 + } + }, + { + "type": "file", + "path": "mmproj-Qwen-Image-2.1-Turbo-Uncensored-F16.gguf", + "size": 1159029984, + "lfs": { + "size": 1159029984 + } + } + ] +} \ No newline at end of file diff --git a/test/campioni-hf/Comfy-Org__vae-text-encorder-for-flux-klein-4b.json b/test/campioni-hf/Comfy-Org__vae-text-encorder-for-flux-klein-4b.json new file mode 100644 index 0000000..77ac689 --- /dev/null +++ b/test/campioni-hf/Comfy-Org__vae-text-encorder-for-flux-klein-4b.json @@ -0,0 +1,56 @@ +{ + "repo": "Comfy-Org/vae-text-encorder-for-flux-klein-4b", + "revisione": "87cfba3d4b98662fb80f23921bb37bf887c97c58", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 1519 + }, + { + "type": "file", + "path": "README.md", + "size": 1737 + }, + { + "type": "file", + "path": "split_files/diffusion_models/flux-2-klein-4b.safetensors", + "size": 7751105712, + "lfs": { + "size": 7751105712 + } + }, + { + "type": "file", + "path": "split_files/diffusion_models/flux-2-klein-base-4b.safetensors", + "size": 7751105712, + "lfs": { + "size": 7751105712 + } + }, + { + "type": "file", + "path": "split_files/text_encoders/qwen_3_4b.safetensors", + "size": 8044982048, + "lfs": { + "size": 8044982048 + } + }, + { + "type": "file", + "path": "split_files/text_encoders/qwen_3_4b_fp4_flux2.safetensors", + "size": 3848213998, + "lfs": { + "size": 3848213998 + } + }, + { + "type": "file", + "path": "split_files/vae/flux2-vae.safetensors", + "size": 336211292, + "lfs": { + "size": 336211292 + } + } + ] +} \ No newline at end of file diff --git a/test/campioni-hf/Helsinki-NLP__opus-mt-it-en.json b/test/campioni-hf/Helsinki-NLP__opus-mt-it-en.json new file mode 100644 index 0000000..c1838d3 --- /dev/null +++ b/test/campioni-hf/Helsinki-NLP__opus-mt-it-en.json @@ -0,0 +1,62 @@ +{ + "repo": "Helsinki-NLP/opus-mt-it-en", + "revisione": "main", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 345 + }, + { + "type": "file", + "path": "README.md", + "size": 901 + }, + { + "type": "file", + "path": "config.json", + "size": 1381 + }, + { + "type": "file", + "path": "generation_config.json", + "size": 293 + }, + { + "type": "file", + "path": "pytorch_model.bin", + "size": 343642677, + "lfs": { + "size": 343642677 + } + }, + { + "type": "file", + "path": "source.spm", + "size": 813709 + }, + { + "type": "file", + "path": "target.spm", + "size": 789549 + }, + { + "type": "file", + "path": "tf_model.h5", + "size": 344197768, + "lfs": { + "size": 344197768 + } + }, + { + "type": "file", + "path": "tokenizer_config.json", + "size": 42 + }, + { + "type": "file", + "path": "vocab.json", + "size": 2369833 + } + ] +} \ No newline at end of file diff --git a/test/campioni-hf/ZhengPeng7__BiRefNet.json b/test/campioni-hf/ZhengPeng7__BiRefNet.json new file mode 100644 index 0000000..8f36d74 --- /dev/null +++ b/test/campioni-hf/ZhengPeng7__BiRefNet.json @@ -0,0 +1,54 @@ +{ + "repo": "ZhengPeng7/BiRefNet", + "revisione": "e2bf8e4460fc8fa32bba5ea4d94b3233d367b0e4", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 1519 + }, + { + "type": "file", + "path": ".gitignore", + "size": 1908 + }, + { + "type": "file", + "path": "BiRefNet_config.py", + "size": 298 + }, + { + "type": "file", + "path": "README.md", + "size": 9965 + }, + { + "type": "file", + "path": "birefnet.py", + "size": 91896 + }, + { + "type": "file", + "path": "config.json", + "size": 405 + }, + { + "type": "file", + "path": "handler.py", + "size": 4762 + }, + { + "type": "file", + "path": "model.safetensors", + "size": 444473596, + "lfs": { + "size": 444473596 + } + }, + { + "type": "file", + "path": "requirements.txt", + "size": 149 + } + ] +} \ No newline at end of file diff --git a/test/campioni-hf/ahmed22xa__Huihui-Qwen3-VL-4B-Instruct-abliterated-comfy.json b/test/campioni-hf/ahmed22xa__Huihui-Qwen3-VL-4B-Instruct-abliterated-comfy.json new file mode 100644 index 0000000..ea71d12 --- /dev/null +++ b/test/campioni-hf/ahmed22xa__Huihui-Qwen3-VL-4B-Instruct-abliterated-comfy.json @@ -0,0 +1,40 @@ +{ + "repo": "ahmed22xa/Huihui-Qwen3-VL-4B-Instruct-abliterated-comfy", + "revisione": "6d6fc98f9bfa783dfa4f143804525742cb5dad62", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 1519 + }, + { + "type": "file", + "path": "Huihui-Qwen3-VL-4B-Instruct-abliterated-fp8_scaled.safetensors", + "size": 5242481504, + "lfs": { + "size": 5242481504 + } + }, + { + "type": "file", + "path": "Huihui-Qwen3-VL-4B-Instruct-abliterated-int8_convrot.safetensors", + "size": 5064440625, + "lfs": { + "size": 5064440625 + } + }, + { + "type": "file", + "path": "Huihui-Qwen3-VL-4B-Instruct-abliterated.safetensors", + "size": 8875719408, + "lfs": { + "size": 8875719408 + } + }, + { + "type": "file", + "path": "README.md", + "size": 7839 + } + ] +} \ No newline at end of file diff --git a/test/campioni-hf/black-forest-labs__FLUX.2-klein-4b-fp8.json b/test/campioni-hf/black-forest-labs__FLUX.2-klein-4b-fp8.json new file mode 100644 index 0000000..dc6f7c5 --- /dev/null +++ b/test/campioni-hf/black-forest-labs__FLUX.2-klein-4b-fp8.json @@ -0,0 +1,53 @@ +{ + "repo": "black-forest-labs/FLUX.2-klein-4b-fp8", + "revisione": "5b4408e59397a4a37ccb46afe426d8ed86379441", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 1662 + }, + { + "type": "file", + "path": "LICENSE.md", + "size": 9584 + }, + { + "type": "file", + "path": "README.md", + "size": 7778 + }, + { + "type": "file", + "path": "editing.jpg", + "size": 2506178, + "lfs": { + "size": 2506178 + } + }, + { + "type": "file", + "path": "flux-2-klein-4b-fp8.safetensors", + "size": 4070624520, + "lfs": { + "size": 4070624520 + } + }, + { + "type": "file", + "path": "others.jpg", + "size": 3386946, + "lfs": { + "size": 3386946 + } + }, + { + "type": "file", + "path": "realism.jpg", + "size": 2855711, + "lfs": { + "size": 2855711 + } + } + ] +} \ No newline at end of file diff --git a/test/campioni-hf/black-forest-labs__flux-3-image-klein-comfy.json b/test/campioni-hf/black-forest-labs__flux-3-image-klein-comfy.json new file mode 100644 index 0000000..ef1167a --- /dev/null +++ b/test/campioni-hf/black-forest-labs__flux-3-image-klein-comfy.json @@ -0,0 +1,240 @@ +{ + "repo": "black-forest-labs/flux-3-image-klein-comfy", + "revisione": "feef60b0aaf63a4773ace96636ebbcf43b881312", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 1561 + }, + { + "type": "file", + "path": "LICENSE.md", + "size": 19593 + }, + { + "type": "file", + "path": "NOTICE.md", + "size": 330 + }, + { + "type": "file", + "path": "README.md", + "size": 12005 + }, + { + "type": "file", + "path": "diffusion_models/flux3-image-klein-base-bf16.safetensors", + "size": 11076957072, + "lfs": { + "size": 11076957072 + } + }, + { + "type": "file", + "path": "diffusion_models/flux3-image-klein-base-int8-convrot.safetensors", + "size": 5807427339, + "lfs": { + "size": 5807427339 + } + }, + { + "type": "file", + "path": "diffusion_models/flux3-image-klein-base-nvfp4.safetensors", + "size": 5257238192, + "lfs": { + "size": 5257238192 + } + }, + { + "type": "file", + "path": "diffusion_models/flux3-image-klein-bf16.safetensors", + "size": 11097404512, + "lfs": { + "size": 11097404512 + } + }, + { + "type": "file", + "path": "diffusion_models/flux3-image-klein-int8-convrot.safetensors", + "size": 5827874779, + "lfs": { + "size": 5827874779 + } + }, + { + "type": "file", + "path": "diffusion_models/flux3-image-klein-nvfp4.safetensors", + "size": 5277685632, + "lfs": { + "size": 5277685632 + } + }, + { + "type": "file", + "path": "launch-gallery.jpg", + "size": 6537996, + "lfs": { + "size": 6537996 + } + }, + { + "type": "file", + "path": "text_encoders/qwen3.5_9b_flux3_prompt_upsampler_bf16.safetensors", + "size": 18819722392, + "lfs": { + "size": 18819722392 + } + }, + { + "type": "file", + "path": "text_encoders/qwen3.5_9b_flux3_prompt_upsampler_int8_convrot.safetensors", + "size": 9748113996, + "lfs": { + "size": 9748113996 + } + }, + { + "type": "file", + "path": "text_encoders/qwen3.5_9b_flux3_prompt_upsampler_w4a8.safetensors", + "size": 6724090780, + "lfs": { + "size": 6724090780 + } + }, + { + "type": "file", + "path": "text_encoders/qwen3vl_4b_bf16.safetensors", + "size": 8875719408, + "lfs": { + "size": 8875719408 + } + }, + { + "type": "file", + "path": "text_encoders/qwen3vl_4b_int8_convrot.safetensors", + "size": 4858559240, + "lfs": { + "size": 4858559240 + } + }, + { + "type": "file", + "path": "vae/flux2-vae.safetensors", + "size": 336211292, + "lfs": { + "size": 336211292 + } + }, + { + "type": "file", + "path": "workflows/LICENSE.upstream", + "size": 1074 + }, + { + "type": "file", + "path": "workflows/README.md", + "size": 1433 + }, + { + "type": "file", + "path": "workflows/flux3_klein_base_bf16_i2i.json", + "size": 109041 + }, + { + "type": "file", + "path": "workflows/flux3_klein_base_bf16_mi2i.json", + "size": 116057 + }, + { + "type": "file", + "path": "workflows/flux3_klein_base_bf16_t2i.json", + "size": 84832 + }, + { + "type": "file", + "path": "workflows/flux3_klein_base_int8_i2i.json", + "size": 109227 + }, + { + "type": "file", + "path": "workflows/flux3_klein_base_int8_mi2i.json", + "size": 116243 + }, + { + "type": "file", + "path": "workflows/flux3_klein_base_int8_t2i.json", + "size": 85018 + }, + { + "type": "file", + "path": "workflows/flux3_klein_bf16_i2i.json", + "size": 107949 + }, + { + "type": "file", + "path": "workflows/flux3_klein_bf16_mi2i.json", + "size": 114965 + }, + { + "type": "file", + "path": "workflows/flux3_klein_bf16_t2i.json", + "size": 83740 + }, + { + "type": "file", + "path": "workflows/flux3_klein_int8_i2i.json", + "size": 108135 + }, + { + "type": "file", + "path": "workflows/flux3_klein_int8_mi2i.json", + "size": 115151 + }, + { + "type": "file", + "path": "workflows/flux3_klein_int8_t2i.json", + "size": 83926 + }, + { + "type": "file", + "path": "workflows/klein_ref_autumn_forest.jpg", + "size": 1213947, + "lfs": { + "size": 1213947 + } + }, + { + "type": "file", + "path": "workflows/klein_ref_great_wave.jpg", + "size": 1348397, + "lfs": { + "size": 1348397 + } + }, + { + "type": "file", + "path": "workflows/klein_ref_green_doors.jpg", + "size": 1283633, + "lfs": { + "size": 1283633 + } + }, + { + "type": "file", + "path": "workflows/klein_ref_toast.jpg", + "size": 1166604, + "lfs": { + "size": 1166604 + } + }, + { + "type": "file", + "path": "workflows/klein_ref_woman_hat.jpg", + "size": 589087, + "lfs": { + "size": 589087 + } + } + ] +} \ No newline at end of file diff --git a/test/campioni-hf/facebook__sam2.1-hiera-small.json b/test/campioni-hf/facebook__sam2.1-hiera-small.json new file mode 100644 index 0000000..53e87c9 --- /dev/null +++ b/test/campioni-hf/facebook__sam2.1-hiera-small.json @@ -0,0 +1,57 @@ +{ + "repo": "facebook/sam2.1-hiera-small", + "revisione": "ee5bba1d82bb8749febdf90f45e84b687142ba03", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 1519 + }, + { + "type": "file", + "path": "README.md", + "size": 19974 + }, + { + "type": "file", + "path": "config.json", + "size": 5698 + }, + { + "type": "file", + "path": "model.safetensors", + "size": 184305280, + "lfs": { + "size": 184305280 + } + }, + { + "type": "file", + "path": "preprocessor_config.json", + "size": 683 + }, + { + "type": "file", + "path": "processor_config.json", + "size": 95 + }, + { + "type": "file", + "path": "sam2.1_hiera_s.yaml", + "size": 3761 + }, + { + "type": "file", + "path": "sam2.1_hiera_small.pt", + "size": 184416285, + "lfs": { + "size": 184416285 + } + }, + { + "type": "file", + "path": "video_preprocessor_config.json", + "size": 705 + } + ] +} \ No newline at end of file diff --git a/test/campioni-hf/mradermacher__Qwen3-VL-8B-Instruct-Uncensored-V2-GGUF.json b/test/campioni-hf/mradermacher__Qwen3-VL-8B-Instruct-Uncensored-V2-GGUF.json new file mode 100644 index 0000000..3c2f6e7 --- /dev/null +++ b/test/campioni-hf/mradermacher__Qwen3-VL-8B-Instruct-Uncensored-V2-GGUF.json @@ -0,0 +1,128 @@ +{ + "repo": "mradermacher/Qwen3-VL-8B-Instruct-Uncensored-V2-GGUF", + "revisione": "main", + "file": [ + { + "type": "file", + "path": ".gitattributes", + "size": 2681 + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.IQ4_XS.gguf", + "size": 4593298272, + "lfs": { + "size": 4593298272 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q2_K.gguf", + "size": 3281734496, + "lfs": { + "size": 3281734496 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q3_K_L.gguf", + "size": 4431395680, + "lfs": { + "size": 4431395680 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q3_K_M.gguf", + "size": 4124162912, + "lfs": { + "size": 4124162912 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q3_K_S.gguf", + "size": 3769613152, + "lfs": { + "size": 3769613152 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q4_K_M.gguf", + "size": 5027785568, + "lfs": { + "size": 5027785568 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q4_K_S.gguf", + "size": 4802014048, + "lfs": { + "size": 4802014048 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q5_K_M.gguf", + "size": 5851114336, + "lfs": { + "size": 5851114336 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q5_K_S.gguf", + "size": 5720763232, + "lfs": { + "size": 5720763232 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q6_K.gguf", + "size": 6725901152, + "lfs": { + "size": 6725901152 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.Q8_0.gguf", + "size": 8709520224, + "lfs": { + "size": 8709520224 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.f16.gguf", + "size": 16388045664, + "lfs": { + "size": 16388045664 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.mmproj-Q8_0.gguf", + "size": 752290016, + "lfs": { + "size": 752290016 + } + }, + { + "type": "file", + "path": "Qwen3-VL-8B-Instruct-Uncensored-V2.mmproj-f16.gguf", + "size": 1159030112, + "lfs": { + "size": 1159030112 + } + }, + { + "type": "file", + "path": "README.md", + "size": 4653 + } + ] +} \ No newline at end of file diff --git a/test/catalogo.mjs b/test/catalogo.mjs new file mode 100644 index 0000000..814570a --- /dev/null +++ b/test/catalogo.mjs @@ -0,0 +1,123 @@ +// Test del catalogo (src/main/installa/catalogo.ts + scegli.ts): niente Electron, niente rete (gli alberi dei repo sono in test/campioni-hf/). +// - ogni voce con `repo` trova il suo file nell'albero vero (percorso esatto per Klein, FLUX.2 Klein 4B, BiRefNet, SAM2, LLM uncensored) +// - le voci uncensored hanno il flag, l'avviso e la guardia di scaricamento funziona +// - ogni voce ha licenza e ruolo; i modelli Qwen Research hanno l'avviso; nessuna voce uncensored è `consigliato` (non si scarica mai da sola) +// +// Esecuzione: node test/catalogo.mjs +import { readFileSync, readdirSync } from 'node:fs' +import { dirname, join } from 'node:path' +import { fileURLToPath, pathToFileURL } from 'node:url' + +const QUI = dirname(fileURLToPath(import.meta.url)) +const { CATALOGO, motivoDiBlocco, motivoLicenza, AVVISO_UNCENSORED } = await import(pathToFileURL(join(QUI, '..', 'src', 'main', 'installa', 'catalogo.ts')).href) +const { scegli, tronco, TIPI_SAFETENSORS } = await import(pathToFileURL(join(QUI, '..', 'src', 'main', 'installa', 'scegli.ts')).href) + +let falliti = 0 +let fatti = 0 +function uguale(nome, avuto, atteso) { + fatti++ + const a = JSON.stringify(avuto) + const b = JSON.stringify(atteso) + if (a === b) return + falliti++ + console.log(`NO ${nome}\n avuto: ${a}\n atteso: ${b}`) +} + +const alberi = {} +for (const f of readdirSync(join(QUI, 'campioni-hf'))) { + const j = JSON.parse(readFileSync(join(QUI, 'campioni-hf', f), 'utf8')) + alberi[j.repo] = j +} +const trova = (v) => { + const albero = alberi[v.repo] + if (!albero) return `(manca l'albero di ${v.repo})` + const pesi = TIPI_SAFETENSORS.has(v.tipo) ? /\.(safetensors|bin)$/i : /\.gguf$/i + const f = scegli(albero.file, v.cerca, true, pesi) + return f ? f.path : null +} +const voce = (id) => CATALOGO.find((v) => v.id === id) + +// Klein 3: percorsi esatti dell'albero del 22/10/2026 (e non il ripiego largo) +uguale('klein-int8', trova(voce('klein-int8')), 'diffusion_models/flux3-image-klein-int8-convrot.safetensors') +uguale('klein-bf16', trova(voce('klein-bf16')), 'diffusion_models/flux3-image-klein-bf16.safetensors') +uguale('klein-base-int8', trova(voce('klein-base-int8')), 'diffusion_models/flux3-image-klein-base-int8-convrot.safetensors') +uguale('klein-base-bf16', trova(voce('klein-base-bf16')), 'diffusion_models/flux3-image-klein-base-bf16.safetensors') +uguale('klein-te (non il PUP)', trova(voce('klein-te')), 'text_encoders/qwen3vl_4b_int8_convrot.safetensors') +uguale('klein-vae', trova(voce('klein-vae')), 'vae/flux2-vae.safetensors') +// i ripieghi larghi, da soli, non devono mai prendere l'upsampler (qwen3.5 9b) né il Base per il veloce +for (const id of ['klein-int8', 'klein-bf16']) { + const v = voce(id) + const larga = scegli(alberi[v.repo].file, [v.cerca[v.cerca.length - 1]], true, /\.(safetensors|bin)$/i) + uguale(`ripiego largo di ${id}: non Base, non upsampler`, !!larga && !/base|upsampl/i.test(larga.path), true) +} +for (const id of ['klein-base-int8', 'klein-base-bf16']) { + const v = voce(id) + const larga = scegli(alberi[v.repo].file, [v.cerca[v.cerca.length - 1]], true, /\.(safetensors|bin)$/i) + uguale(`ripiego largo di ${id}: è Base`, !!larga && /base/i.test(larga.path), true) +} +// i nomi vecchi (flux3_image_klein_sd_int8_convrot) restano riconosciuti +{ + const v = voce('klein-int8') + const vecchi = [{ type: 'file', path: 'flux3_image_klein_sd_int8_convrot.safetensors' }, { type: 'file', path: 'flux3_image_klein_base_int8_convrot.safetensors' }] + uguale('nome vecchio klein-int8', scegli(vecchi, v.cerca, true, /\.safetensors$/i)?.path, 'flux3_image_klein_sd_int8_convrot.safetensors') + uguale('nome vecchio klein-base-int8', scegli(vecchi, voce('klein-base-int8').cerca, true, /\.safetensors$/i)?.path, 'flux3_image_klein_base_int8_convrot.safetensors') +} + +// nuove voci +uguale('FLUX.2 Klein 4B fp8', trova(voce('klein4b-fp8')), 'flux-2-klein-4b-fp8.safetensors') +uguale('FLUX.2 Klein 4B encoder', trova(voce('klein4b-te')), 'split_files/text_encoders/qwen_3_4b.safetensors') +uguale('BiRefNet', trova(voce('sfondo-birefnet')), 'model.safetensors') +uguale('SAM2 small', trova(voce('seg-sam2')), 'model.safetensors') +uguale('LLM uncensored 4B', trova(voce('llm-huihui-4b')), 'Huihui-Qwen3-VL-4B-Instruct-abliterated-int8_convrot.safetensors') +// i file accanto (config, tokenizer) ci sono nell'albero +for (const id of ['sfondo-birefnet', 'seg-sam2']) { + const v = voce(id) + const nomi = alberi[v.repo].file.map((f) => tronco(f.path)) + for (const re of v.accanto) uguale(`${id}: accanto ${re}`, nomi.some((n) => new RegExp(re, 'i').test(n)), true) +} +// i repo con revisione fissata: l'albero è stato salvato a quella revisione +for (const v of CATALOGO.filter((x) => x.repo && x.revisione)) uguale(`revisione di ${v.id}`, alberi[v.repo].revisione, v.revisione) +for (const id of ['klein-int8', 'klein-te', 'klein4b-fp8', 'klein4b-te', 'sfondo-birefnet', 'seg-sam2', 'llm-huihui-4b']) uguale(`${id} ha una revisione fissata`, /^[0-9a-f]{40}$/.test(voce(id).revisione || ''), true) + +// uncensored: flag, avviso e guardia +const unc = CATALOGO.filter((v) => v.uncensored).map((v) => v.id).sort() +uguale('le voci uncensored', unc, ['llm-huihui-4b', 'lora-uc', 'te-ac-q4', 'te-ac-q6', 'te-ac-q8', 'te-mr-q4', 'te-mr-q6', 'te-mr-q8']) +for (const v of CATALOGO.filter((x) => x.uncensored)) { + uguale(`${v.id}: avviso`, (v.avviso || "").endsWith(AVVISO_UNCENSORED), true) + uguale(`${v.id}: non consigliato (mai scaricato da solo)`, !!v.consigliato, false) + uguale(`${v.id}: guardia con interruttore spento`, /Modelli uncensored/.test(motivoDiBlocco(v, { uncensored: false, uncensoredAccettato: true }) || ''), true) + uguale(`${v.id}: guardia senza accettazione`, /accetta/.test(motivoDiBlocco(v, { uncensored: true, uncensoredAccettato: false }) || ''), true) + uguale(`${v.id}: via libera con tutto acceso`, motivoDiBlocco(v, { uncensored: true, uncensoredAccettato: true }), null) +} +// il predefinito di fabbrica è tutto spento +uguale('guardia con impostazioni vuote', !!motivoDiBlocco(voce('te-ac-q4'), {}), true) +// le voci normali non sono mai bloccate +for (const v of CATALOGO.filter((x) => !x.uncensored)) uguale(`${v.id}: mai bloccata`, motivoDiBlocco(v, { uncensored: false, uncensoredAccettato: false }), null) +// niente voce uncensored fra i denoiser di Klein (licenza Kommunity §4(d): niente scorciatoie per togliere le protezioni) +uguale('nessun denoiser Klein uncensored', CATALOGO.filter((v) => (v.tipo === 'klein' || v.tipo === 'diffusione') && v.uncensored).length, 0) + +// licenza e ruolo su tutte le voci dei modelli (non sui LoRA di terzi senza dato) +for (const v of CATALOGO.filter((x) => ['diffusione', 'klein', 'klein-te', 'klein-vae', 'sfondo', 'segmentatore', 'llm'].includes(x.tipo))) { + uguale(`${v.id}: licenza`, typeof v.licenza === 'string' && v.licenza.length > 0, true) + uguale(`${v.id}: ruolo`, typeof v.ruolo === 'string', true) +} +for (const v of CATALOGO.filter((x) => x.licenza === 'Qwen Research' && x.ruolo === 'denoiser')) uguale(`${v.id}: avviso licenza Qwen`, /Qwen Research/.test(v.avviso || ''), true) +for (const v of CATALOGO.filter((x) => x.licenza === 'FLUX Kommunity')) uguale(`${v.id}: avviso Kommunity`, /Kommunity/.test(v.avviso || ''), true) +// id unici +uguale('id unici', new Set(CATALOGO.map((v) => v.id)).size, CATALOGO.length) + +// licenza valida: tutte le voci uncensored del catalogo la hanno; una senza licenza, con Kommunity o con «non-commercial» è rifiutata anche a interruttori accesi +const TUTTO = { uncensored: true, uncensoredAccettato: true } +for (const v of CATALOGO.filter((x) => x.uncensored)) uguale(`${v.id}: licenza valida`, motivoLicenza(v), null) +const finta = { id: 'x', nome: 'Finto', uncensored: true, licenza: 'Apache-2.0', licenzaUrl: 'https://x' } +uguale('licenza Apache ok', motivoDiBlocco(finta, TUTTO), null) +uguale('senza licenza → no', /licenza dichiarata/.test(motivoDiBlocco({ ...finta, licenza: undefined }, TUTTO) || ''), true) +uguale('senza link alla licenza → no', /licenza dichiarata/.test(motivoDiBlocco({ ...finta, licenzaUrl: undefined }, TUTTO) || ''), true) +uguale('FLUX Kommunity → no', /Kommunity/.test(motivoDiBlocco({ ...finta, licenza: 'FLUX Kommunity' }, TUTTO) || ''), true) +uguale('non-commercial → no', /non permette/.test(motivoDiBlocco({ ...finta, licenza: 'CC-BY-NC-4.0' }, TUTTO) || ''), true) +uguale('Qwen Research senza avviso → no', /avviso/.test(motivoDiBlocco({ ...finta, licenza: 'Qwen Research' }, TUTTO) || ''), true) +uguale('Qwen Research con avviso → ok', motivoDiBlocco({ ...finta, licenza: 'Qwen Research', avviso: 'licenza «Qwen Research»: solo ricerca' }, TUTTO), null) +uguale('voce non uncensored: la licenza non blocca', motivoDiBlocco({ ...finta, uncensored: false, licenza: 'FLUX Kommunity' }, {}), null) + +console.log(falliti ? `\n${falliti} su ${fatti} controlli falliti` : `ok: ${fatti} controlli`) +process.exit(falliti ? 1 : 0) diff --git a/test/confronta_grafi.mjs b/test/confronta_grafi.mjs index 01a2f59..6399ff4 100644 --- a/test/confronta_grafi.mjs +++ b/test/confronta_grafi.mjs @@ -5,7 +5,7 @@ // // Esecuzione: node --experimental-strip-types test/confronta_grafi.mjs import { join, dirname } from 'node:path' -import { writeFileSync, mkdirSync } from 'node:fs' +import { writeFileSync, mkdirSync, copyFileSync } from 'node:fs' import { fileURLToPath, pathToFileURL } from 'node:url' import { execFileSync } from 'node:child_process' import { tmpdir } from 'node:os' @@ -23,6 +23,8 @@ function normalizza(sorgente) { "constructor(readonly memoria: Memoria = 'normale') {}", "memoria: Memoria = 'normale'\n constructor(memoria: Memoria = 'normale') { this.memoria = memoria }" ) + // il regional prompting (regioni.ts) sta in un file a parte: si copia accanto e si importa con l'estensione + s = s.replace("from './regioni'", "from './regioni.ts'") return s } @@ -36,6 +38,7 @@ function estrai(commit) { return percorso } +copyFileSync(join(REPO, 'src', 'main', 'motore', 'regioni.ts'), join(TEMP, 'regioni.ts')) const moduloVecchio = await import(pathToFileURL(estrai(OLD_COMMIT)).href) const moduloNuovo = await import(pathToFileURL(estrai('HEAD')).href) diff --git a/test/finto.js b/test/finto.js index 3655a81..183a723 100644 --- a/test/finto.js +++ b/test/finto.js @@ -33,7 +33,7 @@ impostazioni: { leggi: v(imp), salva: async (m) => Object.assign(imp, m) }, sistema: { controllo: v({}) }, setup: { cerca: v([]), stato: v({ installato: true, motore: true, modelli: [] }), avvia: v(true), annulla: v() }, - modelli: { stato: v([]), file: v({ diffusione: [imp.modello], encoder: [imp.encoder], vae: [imp.vae], upscaler: [imp.upscaler], klein: [], kleinTe: [], kleinVae: [] }), scarica: v({ stato: 'errore', errore: 'finto' }), annulla: v() }, + modelli: { stato: v([]), segmentatori: v([]), inutilizzati: v([]), elimina: v(), file: v({ diffusione: [imp.modello], encoder: [imp.encoder], vae: [imp.vae], upscaler: [imp.upscaler], klein: [], kleinTe: [], kleinVae: [], klein4b: [], klein4bTe: [] }), scarica: v({ stato: 'errore', errore: 'finto' }), annulla: v() }, motore: { info: v({ stato: 'pronto', porta: 8818, vramTotale: 6442450944, vramLibera: 1500000000, gpu: 'NVIDIA GeForce RTX 3060 Laptop GPU', versione: '0.39.0' }), avvia: v(), riavvia: v(), ferma: v(), log: v(['motore pronto']), libera: v(), apriComfy: v() }, lavori: { elenco: v([]), annulla: v(), svuota: v(), pulisci: v(), @@ -50,7 +50,7 @@ }, testo: { scrivi: async (q) => { window.__testo = q; return JSON.stringify([{bbox_2d:[300,380,640,880],label:'red-roofed brown wooden hut',kind:'object'},{bbox_2d:[780,180,940,340],label:'bright yellow sun',kind:'object'},{bbox_2d:[40,500,260,990],label:'woman in a green jacket, standing',kind:'person'}]) } }, galleria: { elenco: v(opere), opera: async (id) => opere.find((o) => o.id === id), importaFoto: v({}), preferita: v(), elimina: v(), mostra: v(), apriCartella: v(), copia: v(), esporta: v(), importa: v({ importate: [], altre: [] }), meta: v(null) }, - lora: { elenco: v([]), aggiorna: v(), importa: v([]), elimina: v(), scarica: v({}), apriCartella: v() }, + lora: { elenco: v([]), compatibili: v([]), aggiorna: v(), importa: v([]), elimina: v(), scarica: v({}), apriCartella: v() }, file: { scegliImmagini: v(['/foto.png']), scegli: v(null), scegliCartella: v(null), salvaTemp: async (dataUrl) => { const p = '/tmp/t' + ++n + '.png'; mappa[p] = dataUrl; return p }, diff --git a/test/grafi_v2.mjs b/test/grafi_v2.mjs new file mode 100644 index 0000000..cd6b086 --- /dev/null +++ b/test/grafi_v2.mjs @@ -0,0 +1,108 @@ +// Test dei grafi dell'engine v2 (src/main/motore/grafi.ts): FLUX.2 Klein 4B, LoRA multipli con pesi, senza sfondo, estrai fedele (SAM 2), riquadri nel prompt. +// grafi.ts importa altri file senza estensione: si impacchetta con esbuild in una cartella temporanea. +// +// Esecuzione: node test/grafi_v2.mjs +import { fileURLToPath, pathToFileURL } from 'node:url' +import { dirname, join } from 'node:path' +import { mkdtempSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { buildSync } from 'esbuild' + +const QUI = dirname(fileURLToPath(import.meta.url)) +const fuori = join(mkdtempSync(join(tmpdir(), 'daprod-grafi-')), 'grafi.mjs') +buildSync({ + entryPoints: [join(QUI, '..', 'src', 'main', 'motore', 'grafi.ts')], + bundle: true, + format: 'esm', + platform: 'node', + outfile: fuori, + alias: { '@shared': join(QUI, '..', 'src', 'shared') }, + logLevel: 'error' +}) +const G = await import(pathToFileURL(fuori).href) + +let falliti = 0 +let fatti = 0 +function uguale(nome, avuto, atteso) { + fatti++ + const a = JSON.stringify(avuto) + const b = JSON.stringify(atteso) + if (a === b) return + falliti++ + console.log(`NO ${nome}\n avuto: ${a}\n atteso: ${b}`) +} + +const imp = { + cartellaModelli: 'C:/models', cartellaGalleria: 'C:/gallery', + modello: 'qwen-image-2.1-UC-Q4_K_M.gguf', encoder: 'qwen3-vl-8b.safetensors', vae: 'ae.safetensors', upscaler: '4x.safetensors', + famiglia: 'qwen', kleinModello: 'klein4b_fp8.safetensors', kleinBase: '', kleinEncoder: 'qwen_3_4b.safetensors', kleinVae: 'flux2-vae.safetensors', + upscalerRtx: false, cacheKV: 'gpu', cacheTipo: 'default', encoderSuCpu: false, loraTurbo: '', memoria: 'auto', vramMB: 8192, traduci: false, modelloScelto: 'qwen-base' +} +const base = { + prompt: 'a red fox in a forest', negativo: '', seed: 7, passi: 4, cfg: 1, sampler: 'euler', scheduler: 'simple', megapixel: 1, + larghezza: 1024, altezza: 768, immagini: [], lora: [], quante: 1, trasparente: false +} +const tipi = (g) => Object.values(g.prompt).map((n) => n.class_type) +const nodi = (g, tipo) => Object.values(g.prompt).filter((n) => n.class_type === tipo) +const costruisci = (q, i = imp, ing = { immagini: [] }) => G.costruisci({ ...base, ...q }, i, ing, 'prova') + +// --- FLUX.2 Klein 4B: Crea --- +const k4 = costruisci({ modalita: 'crea', famiglia: 'klein4b' }) +uguale('4B: text encoder di tipo flux2', nodi(k4, 'CLIPLoader')[0].inputs.type, 'flux2') +uguale('4B: scaletta Flux2Scheduler con la misura', [nodi(k4, 'Flux2Scheduler')[0].inputs.steps, nodi(k4, 'Flux2Scheduler')[0].inputs.width, nodi(k4, 'Flux2Scheduler')[0].inputs.height], [4, 1024, 768]) +uguale('4B: CFGGuider a 1 + SamplerCustomAdvanced', [nodi(k4, 'CFGGuider')[0].inputs.cfg, nodi(k4, 'SamplerCustomAdvanced').length], [1, 1]) +uguale('4B: latente vuota Flux2', nodi(k4, 'EmptyFlux2LatentImage').length, 1) +uguale('4B: non c\'è la cache KV di Qwen', tipi(k4).includes('QwenImage21Cache'), false) +const k3 = costruisci({ modalita: 'crea', famiglia: 'klein', passi: 8 }) +uguale('Klein 3: text encoder flux3', nodi(k3, 'CLIPLoader')[0].inputs.type, 'flux3') + +// --- LoRA multipli con pesi: nell'ordine, forza 0 saltata, tetto a ±2 --- +const lora = [{ file: 'a.safetensors', forza: 0.6 }, { file: 'b.safetensors', forza: 0 }, { file: 'c.safetensors', forza: 5 }, { file: 'd.safetensors', forza: -0.5 }] +for (const fam of ['qwen', 'klein', 'klein-base', 'klein4b']) { + const g = costruisci({ modalita: 'crea', famiglia: fam, lora }) + uguale(`LoRA ${fam}: tre applicati (il 0 salta)`, nodi(g, 'LoraLoaderModelOnly').map((n) => [n.inputs.lora_name, n.inputs.strength_model]), [['a.safetensors', 0.6], ['c.safetensors', 2], ['d.safetensors', -0.5]]) +} +// in catena: ognuno prende il modello del precedente +{ + const g = costruisci({ modalita: 'crea', famiglia: 'klein4b', lora }) + const ls = Object.entries(g.prompt).filter(([, n]) => n.class_type === 'LoraLoaderModelOnly') + uguale('LoRA in catena', ls[1][1].inputs.model[0], ls[0][0]) +} +uguale('senza LoRA niente nodi LoRA', nodi(costruisci({ modalita: 'crea', famiglia: 'klein' }), 'LoraLoaderModelOnly').length, 0) + +// --- Senza sfondo --- +const ing1 = { immagini: [{ nome: 'foto.png', larghezza: 1000, altezza: 800 }] } +const sf = costruisci({ modalita: 'sfondo', immagini: ['x.png'] }, imp, ing1) +uguale('sfondo: LoadImage → DaProdSfondo → salva, nessun modello che disegna', tipi(sf).sort(), ['DaProdSfondo', 'LoadImage', 'SaveImage'].sort()) +uguale('sfondo: BiRefNet, colori dei bordi puliti', [nodi(sf, 'DaProdSfondo')[0].inputs.modello, nodi(sf, 'DaProdSfondo')[0].inputs.pulisci_bordi], ['birefnet', true]) +// Klein «senza sfondo» in Crea: BiRefNet dopo il VAE +const kt = costruisci({ modalita: 'crea', famiglia: 'klein', trasparente: true, passi: 8 }) +uguale('Klein trasparente: DaProdSfondo dopo il VAE', [tipi(kt).includes('DaProdSfondo'), tipi(kt).includes('SplitImageWithAlpha')], [true, false]) +const qt = costruisci({ modalita: 'crea', famiglia: 'qwen', trasparente: true, passi: 20 }) +uguale('Qwen trasparente: usa l\'alfa del modello, niente BiRefNet', tipi(qt).includes('DaProdSfondo'), false) + +// --- Estrai fedele (SAM 2) --- +const sg = costruisci({ modalita: 'segmenta', immagini: ['x.png'], segmenta: { box: [0.1, 0.2, 0.6, 0.9], punti: [{ x: 0.3, y: 0.4, positivo: true }] } }, imp, ing1) +uguale('segmenta: nodo SAM 2 con riquadro e punti', [nodi(sg, 'DaProdSegmenta').length, JSON.parse(nodi(sg, 'DaProdSegmenta')[0].inputs.box), JSON.parse(nodi(sg, 'DaProdSegmenta')[0].inputs.punti)], [1, [0.1, 0.2, 0.6, 0.9], [{ x: 0.3, y: 0.4, positivo: true }]]) +uguale('segmenta: ritaglia di serie', nodi(sg, 'DaProdSegmenta')[0].inputs.ritaglia, true) +let errore = '' +try { + costruisci({ modalita: 'segmenta', immagini: ['x.png'], segmenta: {} }, imp, ing1) +} catch (e) { + errore = e.message +} +uguale('segmenta: senza riquadro né punti → errore in italiano', /riquadro|punti/.test(errore), true) + +// --- Riquadri (regional prompting) --- +const regioni = [{ id: 'a', box: [0, 0, 0.5, 0.5], prompt: 'a red fox' }, { id: 'b', box: [0.5, 0.5, 1, 1], prompt: 'a blue door' }] +const testoDi = (g) => nodi(g, 'CLIPTextEncode').map((n) => n.inputs.text).concat(nodi(g, 'TextEncodeQwenImage21').map((n) => n.inputs.prompt))[0] || '' +const rk = costruisci({ modalita: 'crea', famiglia: 'klein', passi: 8, regioni }) +uguale('Klein 3 + riquadri: JSON con le coordinate', /"bbox":\[0,0,500,500\]/.test(testoDi(rk)), true) +const rq = costruisci({ modalita: 'crea', famiglia: 'qwen', passi: 20, regioni }) +uguale('Qwen + riquadri: a parole', /Compose the image with these elements/.test(testoDi(rq)) && !testoDi(rq).includes('bbox'), true) +const r4 = costruisci({ modalita: 'crea', famiglia: 'klein4b', regioni }) +uguale('Klein 4B + riquadri: a parole', /Compose the image with these elements/.test(testoDi(r4)), true) +uguale('senza riquadri il prompt è quello', testoDi(costruisci({ modalita: 'crea', famiglia: 'klein', passi: 8 })), base.prompt) + +console.log(falliti ? `${falliti} FALLITI su ${fatti}` : `TUTTI OK (${fatti} controlli)`) +process.exit(falliti ? 1 : 0) diff --git a/test/lora.mjs b/test/lora.mjs new file mode 100644 index 0000000..235e8cd --- /dev/null +++ b/test/lora.mjs @@ -0,0 +1,50 @@ +// Test del riconoscimento dei LoRA dai nomi dei tensori (src/shared/loraBase.ts). +// +// Esecuzione: node test/lora.mjs +import { fileURLToPath, pathToFileURL } from 'node:url' +import { dirname, join } from 'node:path' + +const QUI = dirname(fileURLToPath(import.meta.url)) +const L = await import(pathToFileURL(join(QUI, '..', 'src', 'shared', 'loraBase.ts')).href) + +let falliti = 0 +let fatti = 0 +function uguale(nome, avuto, atteso) { + fatti++ + const a = JSON.stringify(avuto) + const b = JSON.stringify(atteso) + if (a === b) return + falliti++ + console.log(`NO ${nome}\n avuto: ${a}\n atteso: ${b}`) +} +const ris = (c, f) => { + const a = L.analizzaChiavi(c, f) + return [a.compatibile, a.base, a.famiglie] +} + +uguale('vuoto', ris([]), ['?', undefined, []]) +uguale('Qwen 2.1', ris(['transformer_blocks.3.attn.to_q.lora_A.weight', 'transformer_blocks.3.img_mlp.net.0.proj.lora_B.weight']), ['si', 'Qwen-Image 2.1', ['qwen']]) +uguale('Qwen 1.x', ris(['transformer_blocks.0.txt_mlp.net.0.lora_A.weight']), ['no', 'Qwen-Image 1.x', []]) +uguale('SDXL', ris(['lora_unet_down_blocks_0_attentions_0.lora_down.weight']), ['no', 'Stable Diffusion / SDXL', []]) +uguale('Klein 3 (single_blocks qkv)', ris(['diffusion_model.single_blocks.2.qkv.lora_A.weight']), ['si', 'FLUX 3 Klein', ['klein', 'klein-base']]) +uguale('Klein 3 (content_mode_blocks)', ris(['diffusion_model.content_mode_blocks.0.attn.lora_A.weight']), ['si', 'FLUX 3 Klein', ['klein', 'klein-base']]) + +const doppi = (n) => Array.from({ length: n }, (_, i) => `diffusion_model.double_blocks.${i}.img_attn.qkv.lora_A.weight`) +const flux2 = (n) => [...doppi(n), 'diffusion_model.double_blocks.0.img_mlp.0.lora_A.weight', 'diffusion_model.double_blocks.0.txt_mlp.0.lora_A.weight', 'diffusion_model.single_blocks.0.linear1.lora_A.weight'] +uguale('Klein 4B (5 blocchi doppi)', ris(flux2(5)), ['si', 'FLUX.2 Klein 4B', ['klein4b']]) +uguale('Klein 4B con forma 3072', ris(flux2(5), { 'diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight': [16, 3072] }), ['si', 'FLUX.2 Klein 4B', ['klein4b']]) +uguale('FLUX.2 grande (8 doppi)', ris(flux2(8)), ['no', 'FLUX.1 / FLUX.2 grande', []]) +uguale('FLUX.2 9B (forma 4096)', ris(flux2(5), { 'diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight': [16, 4096] }), ['no', 'FLUX.2', []]) +uguale('sconosciuto', ris(['foo.bar.weight']), ['?', undefined, []]) + +// vaBeneCon +uguale('Qwen con Qwen', L.vaBeneCon({ compatibile: 'si', famiglie: ['qwen'] }, 'qwen'), 'si') +uguale('Qwen con Klein', L.vaBeneCon({ compatibile: 'si', famiglie: ['qwen'] }, 'klein'), 'no') +uguale('Klein 3 vale per Base', L.vaBeneCon({ compatibile: 'si', famiglie: ['klein', 'klein-base'] }, 'klein-base'), 'si') +uguale('Klein 3 non per 4B', L.vaBeneCon({ compatibile: 'si', famiglie: ['klein', 'klein-base'] }, 'klein4b'), 'no') +uguale('vecchio LoRA senza famiglie (si) vale', L.vaBeneCon({ compatibile: 'si' }, 'qwen'), 'si') +uguale('vecchio LoRA senza famiglie (no)', L.vaBeneCon({ compatibile: 'no' }, 'qwen'), 'no') +uguale('non si capisce → si lascia provare', L.vaBeneCon({ compatibile: '?' }, 'klein'), '?') + +console.log(falliti ? `${falliti} FALLITI su ${fatti}` : `TUTTI OK (${fatti} controlli)`) +process.exit(falliti ? 1 : 0) diff --git a/test/modelli.mjs b/test/modelli.mjs index 28039c9..397d0d9 100644 --- a/test/modelli.mjs +++ b/test/modelli.mjs @@ -22,8 +22,8 @@ function uguale(nome, avuto, atteso) { } // la mappa -uguale('tre modelli, in ordine', M.ORDINE_MODELLI, ['qwen-base', 'klein', 'klein-base']) -uguale('le chiavi della mappa', Object.keys(M.MODELLI).sort(), ['klein', 'klein-base', 'qwen-base']) +uguale('quattro modelli, in ordine', M.ORDINE_MODELLI, ['qwen-base', 'klein', 'klein-base', 'klein4b']) +uguale('le chiavi della mappa', Object.keys(M.MODELLI).sort(), ['klein', 'klein-base', 'klein4b', 'qwen-base']) uguale('qwen-base', [M.MODELLI['qwen-base'].famiglia, M.MODELLI['qwen-base'].passi, M.MODELLI['qwen-base'].turbo, M.MODELLI['qwen-base'].lora], ['qwen', 40, false, true]) uguale('klein', [M.MODELLI.klein.famiglia, M.MODELLI.klein.passi, M.MODELLI.klein.cfg, M.MODELLI.klein.veloce], ['klein', 8, 1, true]) uguale('klein-base', [M.MODELLI['klein-base'].famiglia, M.MODELLI['klein-base'].passi, M.MODELLI['klein-base'].cfg], ['klein-base', 50, 3.5]) diff --git a/test/regioni.mjs b/test/regioni.mjs new file mode 100644 index 0000000..caa9c05 --- /dev/null +++ b/test/regioni.mjs @@ -0,0 +1,64 @@ +// Test del regional prompting (src/main/motore/regioni.ts): riquadri → testo per Klein 3 (JSON con le coordinate) e per Qwen/Klein 4B (a parole). +// +// Esecuzione: node test/regioni.mjs +import { fileURLToPath, pathToFileURL } from 'node:url' +import { dirname, join } from 'node:path' + +const QUI = dirname(fileURLToPath(import.meta.url)) +const R = await import(pathToFileURL(join(QUI, '..', 'src', 'main', 'motore', 'regioni.ts')).href) + +let falliti = 0 +let fatti = 0 +function uguale(nome, avuto, atteso) { + fatti++ + const a = JSON.stringify(avuto) + const b = JSON.stringify(atteso) + if (a === b) return + falliti++ + console.log(`NO ${nome}\n avuto: ${a}\n atteso: ${b}`) +} + +// pulizia +uguale('niente regioni', R.regioniPronte(undefined), []) +uguale('prompt vuoto scartato', R.regioniPronte([{ id: 'a', box: [0, 0, 0.5, 0.5], prompt: ' ' }]), []) +uguale('riquadro troppo piccolo scartato', R.regioniPronte([{ id: 'a', box: [0.1, 0.1, 0.11, 0.5], prompt: 'cat' }]), []) +uguale('riquadro con 3 numeri scartato', R.regioniPronte([{ id: 'a', box: [0, 0, 1], prompt: 'cat' }]), []) +uguale('NaN scartato', R.regioniPronte([{ id: 'a', box: [0, 0, NaN, 1], prompt: 'cat' }]), []) +uguale('riquadro capovolto e fuori tela', R.regioniPronte([{ id: 'x', box: [1.2, 0.8, 0.5, -0.2], prompt: ' a red\n ball ' }]), [{ id: 'obj_1', box: [0.5, 0, 1, 0.8], prompt: 'a red ball' }]) +uguale('id obj_N in fila, saltando gli scartati', R.regioniPronte([ + { id: 'a', box: [0, 0, 0.5, 0.5], prompt: 'one' }, + { id: 'b', box: [0, 0, 0.5, 0.5], prompt: '' }, + { id: 'c', box: [0.5, 0.5, 1, 1], prompt: 'two' } +]).map((r) => r.id), ['obj_1', 'obj_2']) +const venti = Array.from({ length: 20 }, (_, i) => ({ id: `r${i}`, box: [0, 0, 0.5, 0.5], prompt: `p${i}` })) +uguale('al massimo 16', R.regioniPronte(venti).length, 16) + +// Klein 3: bbox [y_min, x_min, y_max, x_max] su 0–1000 +uguale('bbox di Klein', R.bboxKlein([0.1, 0.2, 0.6, 0.9]), [200, 100, 900, 600]) +const reg = R.regioniPronte([ + { id: 'a', box: [0, 0, 0.5, 0.5], prompt: 'a red fox' }, + { id: 'b', box: [0.5, 0.5, 1, 1], prompt: 'a blue door.' } +]) +const pk = R.promptKleinConBox('a forest at dawn', reg) +uguale('Klein: didascalia con gli id', pk.startsWith('A forest at dawn. A red fox . A blue door . ['), true) +const json = JSON.parse(pk.slice(pk.indexOf('[{'))) +uguale('Klein: righe JSON', json, [ + { id: 'obj_1', bbox: [0, 0, 500, 500], desc: 'A red fox' }, + { id: 'obj_2', bbox: [500, 500, 1000, 1000], desc: 'A blue door.' } +]) +const pm = JSON.parse(R.promptKleinConBox('', reg, true).slice(R.promptKleinConBox('', reg, true).indexOf('[{'))) +uguale('Klein modifica: tgt_bbox', pm[0], { id: 'obj_1', from: 'target', src_bbox: null, tgt_bbox: [0, 0, 500, 500], desc: 'A red fox' }) +uguale('senza regioni il testo non cambia', R.promptKleinConBox('x', []), 'x') + +// Qwen / Klein 4B: a parole +uguale('dove: in alto a sinistra, piccolo', R.dove([0.02, 0.02, 0.2, 0.2]), 'small, at the top left of the frame') +uguale('dove: al centro', R.dove([0.35, 0.35, 0.65, 0.65]), 'in the center of the frame') +uguale('dove: grande a destra', R.dove([0.5, 0.1, 1, 0.9]), 'large, at the right of the frame') +const pq = R.promptConPosizioni('a street', reg) +uguale('Qwen: scena + elementi', pq.startsWith('A street. Compose the image with these elements, each in its place:'), true) +uguale('Qwen: nessun JSON', pq.includes('{'), false) +uguale('Qwen modifica', R.promptConPosizioni('', reg, true).startsWith('Apply these changes, each one in its place:'), true) +uguale('senza regioni (Qwen)', R.promptConPosizioni('x', []), 'x') + +console.log(falliti ? `${falliti} FALLITI su ${fatti}` : `TUTTI OK (${fatti} controlli)`) +process.exit(falliti ? 1 : 0) diff --git a/test/spazio.mjs b/test/spazio.mjs index 1b41dce..52b766f 100644 --- a/test/spazio.mjs +++ b/test/spazio.mjs @@ -23,6 +23,9 @@ function uguale(nome, avuto, atteso) { uguale('qwen', S.occupaImmagine('qwen'), 'immagine:qwen') uguale('klein', S.occupaImmagine('klein'), 'immagine:klein') uguale('klein-base', S.occupaImmagine('klein-base'), 'immagine:klein-base') +uguale('klein4b', S.occupaImmagine('klein4b'), 'immagine:klein4b') +uguale('da Klein 4B a Qwen si libera', S.cambioDiModello('immagine:klein4b', 'immagine:qwen').libera, true) +uguale('fase di scarico Klein 4B', S.faseScarico('immagine:klein4b'), 'Scarico FLUX.2 Klein 4B') uguale('testo col 4B', S.occupaTesto(true), 'testo:4b') uguale('testo con l\'8B', S.occupaTesto(false), 'testo:8b') uguale('testo senza dire niente = 8B', S.occupaTesto(undefined), 'testo:8b')