UniTune: Googles alternativa för neuralt bildredigering

Google Research verkar attackera textbaserad bildredigering från flera håll och väntar förmodligen på att se vad som “fungerar”. Strax efter att de släppte sin Imagic-rapport föreslog sökjätten en ytterligare latent diffusionsbaserad metod för att utföra tidigare omöjliga AI-baserade redigeringar av bilder via textkommandon, den här gången kallad UniTune.

Utifrån de exempel som ges i projektets nya rapport har UniTune uppnått en extraordinär grad av disentanglement av semantisk pose och idé från faktiskt innehåll i bilden:

UniTunes kommando över semantisk komposition är utmärkt. Observera hur ansiktena på de två personerna i den översta raden av bilder inte har förvrängts av den extraordinära transformationen av resten av källbilden (höger). Källa: https://arxiv.org/pdf/2210.09477.pdf

Som fans av Stable Diffusion har lärt sig kan det vara svårt att applicera redigeringar på delar av en bild utan att negativt påverka resten av bilden. Även om populära distributioner som AUTOMATIC1111 kan skapa masker för lokala och begränsade redigeringar, är processen torterande och ofta oförutsägbar.

Det uppenbara svaret, åtminstone för en datorseendepraktiker, är att införa ett lager av semantisk segmentering som kan känna igen och isolera objekt i en bild utan användarintervention, och faktum är att det har funnits flera nya initiativ på sistone i den här riktningen.

En annan möjlighet för att låsa fast oordnade och sammanflätade neuralt bildredigeringsoperationer är att utnyttja OpenAI:s inflytelserika Contrastive Language–Image Pre-training (CLIP) modul, som ligger till grund för latent diffusionsmodeller som DALL-E 2 och Stable Diffusion, för att fungera som en filter vid den punkt då en text-till-bild-modell är redo att skicka en tolkad render tillbaka till användaren. I detta sammanhang bör CLIP fungera som en väktare och kvalitetskontrollmodul, som avvisar felaktiga eller olämpliga render. Detta är på väg att införas (Discord-länk) på Stability.ai:s DreamStudio API-driven portal.

Men eftersom CLIP kan ses som både boven och lösningen i ett sådant scenario (eftersom det också informerade om hur bilden utvecklades), och eftersom maskinvarukraven kan överstiga vad som sannolikt finns tillgängligt lokalt för en slutanvändare, kan denna metod inte vara idealisk.

Komprimerat språk

Den föreslagna UniTune “finjusterar” en befintlig diffusionsmodell – i det här fallet Googles egna Imagen, även om forskarna påstår att metoden är kompatibel med andra latent diffusionsarkitekturer – så att en unik token injiceras i den som kan anropas genom att inkludera den i en textprompt.

Utifrån vad som kan ses verkar detta som Googles DreamBooth, som för närvarande är en besatthet bland Stable Diffusion-fans och utvecklare, som kan injicera nya karaktärer eller objekt i en befintlig checkpoint, ofta på mindre än en timme, baserat på ett fåtal källbilder; eller också som Textual Inversion, som skapar “sidecar”-filer för en checkpoint, som sedan behandlas som om de ursprungligen tränades in i modellen, och kan dra nytta av modellens egna omfattande resurser genom att modifiera dess textklassificerare, vilket resulterar i en liten fil (jämfört med de minst 2 GB pruned checkpoints som DreamBooth har).

I själva verket hävdar forskarna att UniTune avvisade båda dessa tillvägagångssätt. De fann att Textual Inversion utelämnade för många viktiga detaljer, medan DreamBooth “presterade sämre och tog längre tid” än den lösning de slutligen fastnade för.

Trots allt använder UniTune samma inkapslade semantiska “metaprompt”-tillvägagångssätt som DreamBooth, med tränade förändringar som anropas av unika ord valda av tränaren, som inte kommer att krocka med några termer som redan finns i en noggrant tränad offentlig modell.

‘För att utföra redigeringsoperationen, sampar vi de finjusterade modellerna med prompten “[sällsynta_token] edit_prompt” (t.ex. “beikkpic två hundar i en restaurang” eller “beikkpic en minion”).’

Processen

Även om det är förvirrande varför två nästan identiska rapporter, i termer av deras slutliga funktion, ska anlända från Google under samma vecka, finns det, trots en stor mängd likheter mellan de två initiativen, åtminstone en tydlig skillnad mellan UniTune och Imagic – den senare använder “okomprimerat” naturligt språk för att vägleda bildredigeringsoperationer, medan UniTune tränar in unika DreamBooth-liknande token.

Därför, om du redigerade med Imagic och ville åstadkomma en transformation av den här naturen…

Från UniTune-rapporten – UniTune ställer sig mot Googles favoritkonkurrent neuralt redigeringsramverk, SDEdit. UniTunes resultat är på den yttersta högersidan, medan den uppskattade masken ses i den andra bilden från vänster.

… i Imagic skulle du mata in ‘den tredje personen, sittande i bakgrunden, som en söt pälsmonster’.

Det motsvarande UniTune-kommandot skulle vara ‘Killen bakom som [x]’, där x är vilket konstigt och unikt ord som var bundet till den finjusterade konceptet associerat med den pälsmonstertypen.

Medan flera bilder matas in i antingen DreamBooth eller Textual Inversion med avsikt att skapa en deepfake-liknande abstraktion som kan anropas i många poser, matar både UniTune och Imagic en enda bild in i systemet – den ursprungliga, oförändrade bilden.

Detta liknar sättet som många GAN-baserade redigeringsverktyg har fungerat under de senaste åren – genom att omvandla en inmatad bild till latenta koder i GAN:s latenta utrymme och sedan adressera dessa koder och skicka dem till andra delar av det latenta utrymmet för modifiering (dvs. inmatning av en bild av en ung mörkhårig person och projicering genom latenta koder associerade med “gammal” eller “blond”, etc.).

Men resultaten, i en diffusionsmodell och med denna metod, är ganska anmärkningsvärt exakta i jämförelse, och mycket mindre tvetydiga:

Finjusteringsprocessen

UniTune-metoden skickar i princip den ursprungliga bilden genom en diffusionsmodell med en uppsättning instruktioner om hur den ska modifieras, med hjälp av de stora tillgängliga dataresurserna som tränats in i modellen. I själva verket kan du göra detta nu med Stable Diffusions img2img-funktionalitet – men inte utan att förvränga eller på något sätt ändra delar av bilden som du helst vill behålla.

Under UniTune-processen justeras systemet, det vill säga att UniTune tvingar modellen att återuppta träningen, med de flesta av dess lager otillfrusna (se nedan). I de flesta fall kommer finjustering att sänka de övergripande allmänna förlustvärdena för en hårt vunnen högpresterande modell till förmån för att injicera eller förbättra någon annan aspekt som önskas skapas eller förbättras.

Men med UniTune verkar det som att modellkopior som påverkas, även om de kan väga flera gigabyte eller mer, kommer att behandlas som en engångs “skal” och kasseras i slutet av processen, efter att ha tjänat ett enda syfte. Den här sortens nonchalant datahantering blir en vardaglig lagringskris för DreamBooth-fans, vars egna modeller, även när de är pruned, inte är mindre än 2 GB per ämne.

Som med Imagic sker den huvudsakliga justeringen i UniTune vid de lägsta två av de tre lagren i Imagen (bas 64px, 64px>256px och 256px>1024px). Till skillnad från Imagic ser forskarna något potentiellt värde i att optimera justeringen också för detta sista och största superupplösningslagret (även om de inte har försökt det ännu).

För det lägsta 64px-lagret är modellen biasad mot basbilden under träning, med flera dubbletter av bild/text-par som matas in i systemet under 128 iterationer med en batchstorlek på 4, och med Adafactor som förlustfunktion, som opererar med en inlärningshastighet på 0,0001. Även om T5-kodaren ensam är frusen under denna finjustering, är den också frusen under primärträning av Imagen

Den ovanstående operationen upprepas sedan för 64>256px-lagret, med samma brusförstärkningsförfarande som används i den ursprungliga träningen av Imagen.

Sampling

Det finns många möjliga sampelmetoder som kan användas för att framkalla de förändringar som gjorts, inklusive Classifier Free Guidance (CFG), en huvudfunktion i Stable Diffusion. CFG definierar i princip utsträckningen av vilken modellen är fri att “följa sin fantasi” och utforska renderingmöjligheterna – eller också, vid lägre inställningar, utsträckningen av vilken den bör följa indatakällan och göra färre svepande eller dramatiska förändringar.

Liksom Textual Inversion (lite mindre med DreamBooth) är UniTune mottaglig för att applicera distinkta grafiska stilar på ursprungliga bilder, samt mer fotorealistiska redigeringar.

Forskarna experimenterade också med SDEdits “late start”-teknik, där systemet uppmuntras att bevara originaldetaljer genom att bara delvis vara “brus” från början, men snarare behålla sina väsentliga egenskaper. Även om forskarna bara använde detta på det lägsta lagret (64px), tror de att det kan vara en användbar sampelmetod i framtiden.

Forskarna utnyttjade också prompt-till-prompt som en ytterligare textbaserad teknik för att villkora modellen:

‘I “prompt till prompt”-inställningen fann vi att en teknik som vi kallar Prompt Guidance är särskilt användbar för att justera trohet och uttrycksfullhet.

‘Prompt Guidance är liknande Classifier Free Guidance, men baseline är en annan prompt istället för den ovillkorliga modellen. Detta guider modellen mot delta mellan de två prompterna.’

Prompt-till-prompt i UniTune, som effektivt isolerar områden som ska ändras.

Men prompt-guidance, påstår författarna, behövdes bara ibland i fall där CFG misslyckades med att uppnå det önskade resultatet.

En annan ny sampelmetod som upptäcktes under utvecklingen av UniTune var interpolation, där områden i bilden är tillräckligt distinkta så att både den ursprungliga och den ändrade bilden är mycket lika i komposition, vilket tillåter en mer “naiv” interpolation att användas.

Interpolation kan göra de högre ansträngningarna i UniTune överflödiga i fall där områden som ska transformeras är diskreta och välavgränsade.

Författarna föreslår att interpolation potentiellt kan fungera så bra att den kan användas som standardinställning och observerar också att den har kraften att åstadkomma extraordinära transformationer i fall där komplexa ocklusioner inte behöver förhandlas av mer intensiva metoder.

UniTune kan utföra lokala redigeringar med eller utan redigeringsmasker, men kan också bestämma var redigeringar ska placeras på egen hand, med en ovanlig kombination av tolkningskraft och robust essentialisering av källdata:

I den översta bilden i den andra kolumnen har UniTune, som fått i uppgift att infoga en ‘röd tåg i bakgrunden’, placerat den på en lämplig och autentisk position. Observera i de andra exemplen hur semantisk integritet till källbilden upprätthålls, även i mitten av extraordinära förändringar i pixelinnehåll och kärnstilar i bilderna.

Latens

Även om den första iterationen av ett nytt system kommer att vara långsam, och även om det är möjligt att antingen communityengagemang eller företagsengagemang (det är vanligtvis inte båda) så småningom kommer att påskynda och optimera en resurskrävande rutin, så är både UniTune och Imagic på väg att utföra ganska stora maskinlärningsmanövrer för att skapa dessa fantastiska redigeringar, och det är tveksamt i vilken utsträckning en sådan resurskrävande process någonsin kan skalas ner till hemmabruk, snarare än API-driven åtkomst (även om den senare kan vara mer önskvärd för Google).

Just nu är rundresan från indata till resultat cirka 3 minuter på en T4-GPU, med cirka 30 sekunder extra för inferens (i enlighet med vilken inferensrutin som helst). Författarna medger att detta är hög latens och knappast kvalificerar sig som “interaktivt”, men de påpekar också att modellen förblir tillgänglig för ytterligare redigeringar en gång den har justerats, tills användaren är klar med processen, vilket minskar redigerings­tiden per redigering.

Publicerad första gången den 21 oktober 2022.