Teaser and more readme

README.md

@@ -1,10 +1,39 @@
-# cdim
-Constrained Diffusion Implicit Models
+# Constrained Diffusion Implicit Models
+![alt text](Teaser.jpg)
 
+## Authors
+[Vivek Jayaram](http://www.vivekjayaram.com/), [John Thickstun](https://johnthickstun.com/), [Ira Kemelmacher-Shlizerman](https://homes.cs.washington.edu/~kemelmi/), and [Steve Seitz](https://homes.cs.washington.edu/~seitz/)
+
+## [Project Page](www.google.com)
+(Coming Soon)
+
+### [Paper](www.google.com)
+(Coming Soon)
+
+### Summary
+We solve noisy linear inverse problems with diffusion models. The method is fast and addresses many problems like inpainting, super-resolution, gaussian deblur, and poisson noise. 
+
+
+## Getting started 
+
+### 1) Clone the repository
+
+```
+git clone https://github.com/vivjay30/cdim
+
+cd cdim
+
+export PYTHONPATH=$PYTHONPATH:`pwd`
+```
+
+### 2) Install dependencies
+
+```
 conda create -n cdim python=3.11
 
 conda activate cdim
 
 pip install -r requirements.txt
 
-pip install torch==2.4.1+cu124 torchvision-0.19.1+cu124 --extra-index-url https://download.pytorch.org/whl/cu124
+pip install torch==2.4.1+cu124 torchvision-0.19.1+cu124 --extra-index-url https://download.pytorch.org/whl/cu124
+```

cdim/etas.json

@@ -1,6 +1,6 @@
 {
     "box_inpainting": {
-        "kl" : {
+        "kl": {
             "T50_K3": {
                 "lambda": 4.0,
                 "etas": {"20": 0.0007644674769138669, "40": 0.0007352405159032115, "60": 0.0007339323611117017, "80": 0.0007407785604665967, "100": 0.0007549588527783352, "120": 0.0007624788939223991, "140": 0.0007797297457310411, "160": 0.0008114829360981001, "180": 0.0008749328277143722, "200": 0.0009383883323579534, "220": 0.0010238270356100192, "240": 0.0010784138469520864, "260": 0.0011866835135642144, "280": 0.0013024192080464585, "300": 0.001475853782213996, "320": 0.001632385599489279, "340": 0.0018509140273253787, "360": 0.0020580830457055752, "380": 0.0023878463159672523, "400": 0.0026761896618615814, "420": 0.0031795443584943708, "440": 0.0034981176066313513, "460": 0.004353874833222008, "480": 0.004711352191433642, "500": 0.005726754667302881, "520": 0.006623946294438855, "540": 0.00863887341585953, "560": 0.009790771703642386, "580": 0.011950644941955498, "600": 0.013428242186066879, "620": 0.017515641845500456, "640": 0.01918877035399789, "660": 0.024788169966172627, "680": 0.026881696669367267, "700": 0.034675169438193654, "720": 0.039756030503755736, "740": 0.04854181173586306, "760": 0.05583204092049454, "780": 0.06820393054365007, "800": 0.0775198126468748, "820": 0.09482237250830504, "840": 0.10625708906759426, "860": 0.1277688814432282, "880": 0.1429259541984142, "900": 0.1730718536839941, "920": 0.19272152879630572, "940": 0.2197799661143284, "960": 0.2523353552270739, "980": 0.24785875279032762}
@@ -9,10 +9,20 @@
                 "lambda": 9.0,
                 "etas": {"40": 0.002007311321255281, "80": 0.0017914309657763605, "120": 0.0018785649978129587, "160": 0.0019939875678122057, "200": 0.002077116515854901, "240": 0.0022416421013690723, "280": 0.002525056453252015, "320": 0.002990539520210802, "360": 0.003687435867772619, "400": 0.00461278213332024, "440": 0.005670295237520474, "480": 0.007450458259252178, "520": 0.01104231419851132, "560": 0.01304287128591384, "600": 0.01770141680890715, "640": 0.02673789939069817, "680": 0.03434346958799475, "720": 0.04184474771847407, "760": 0.05826798834659389, "800": 0.07392334565984734, "840": 0.09130451088376304, "880": 0.11364396078616029, "920": 0.1595581552644202, "960": 0.17525405659705573}
             }
+        },
+        "l2": {
+            "T50_K3": {
+                "lambda": 8.0,
+                "etas": {"20": 0.00021444517006169746, "40": 0.00011604119540597255, "60": 9.58772243330201e-05, "80": 8.975413448366393e-05, "100": 9.411635656111335e-05, "120": 8.72203542480862e-05, "140": 9.063772721288521e-05, "160": 9.453677356617138e-05, "180": 0.00012265401768830916, "200": 0.000170195596865024, "220": 0.00029837127400063557, "240": 0.0004947990749204956, "260": 0.0007326859870617194, "280": 0.0009746337397007528, "300": 0.0012234155451859986, "320": 0.0015186419480038773, "340": 0.0018382499373895186, "360": 0.002217853150885955, "380": 0.0026526054340428235, "400": 0.0030888360650993016, "420": 0.003682793771240093, "440": 0.004286529952535807, "460": 0.005192898620704694, "480": 0.0060829834535095365, "500": 0.007351511127457126, "520": 0.00831287559645274, "540": 0.010537466592395335, "560": 0.012053684440492244, "580": 0.014704738301449458, "600": 0.017653333981888054, "620": 0.02158249110142089, "640": 0.02554462233400011, "660": 0.030338915745320535, "680": 0.038534455437019444, "700": 0.04679018227906465, "720": 0.06423099339206706, "740": 0.0713087848211941, "760": 0.09903813634514312, "780": 0.13436430543557223, "800": 0.17160126751236052, "820": 0.24457504945536157, "840": 0.3156838921021669, "860": 0.5075618854250409, "880": 0.6872538380456742, "900": 0.8775862648540197, "920": 1.2991397403837937, "940": 1.5632183751816855, "960": 1.9215644370773675, "980": 1.9215663784225088}
+            },
+            "T25_K1": {
+                "lambda": 10.0,
+                "etas": {"40": 0.00018921975532152269, "80": 0.0001157824796590529, "120": 0.0001153207100068223, "160": 0.00012495355164101671, "200": 0.00013316237708023123, "240": 0.0001548524583230766, "280": 0.00018818787522232978, "320": 0.00023773891755081355, "360": 0.00029604922377423953, "400": 0.0003623565260250032, "440": 0.0004768631986763448, "480": 0.000681236082525412, "520": 0.0008140604233677697, "560": 0.001285574813033369, "600": 0.0016125729795840645, "640": 0.0026998835127931917, "680": 0.0041898900321471035, "720": 0.006548597277460448, "760": 0.009768957330768692, "800": 0.02236172946116906, "840": 0.030298673371512694, "880": 0.0559994584015136, "920": 0.09495997176807491, "960": 0.0795712319137785}
+            }
         }
     },
     "random_inpainting": {
-        "kl" : {
+        "kl": {
             "T50_K3": {
                 "lambda": 3.0,
                 "etas": {"20": 4.2155235275353885e-05, "40": 2.9131803821087877e-05, "60": 3.138850996275609e-05, "80": 3.386381497611811e-05, "100": 3.631776565956598e-05, "120": 3.973121110139383e-05, "140": 4.3322363720455985e-05, "160": 4.863932933274831e-05, "180": 5.240606242179506e-05, "200": 5.814939458443417e-05, "220": 6.197553145350107e-05, "240": 6.662394612209601e-05, "260": 7.353187590939113e-05, "280": 7.957012424639822e-05, "300": 8.54105910002775e-05, "320": 9.482279846848642e-05, "340": 0.00010283310990513798, "360": 0.00011288007608863972, "380": 0.00012768665988629105, "400": 0.00014661738441383113, "420": 0.00016301237603343505, "440": 0.00018253602035619805, "460": 0.00020806616934057048, "480": 0.00022939095079289252, "500": 0.00024356857801398671, "520": 0.00027035908261852304, "540": 0.0003041859182782007, "560": 0.0003671558660027532, "580": 0.00043392769502977437, "600": 0.0004871999766289279, "620": 0.0006560545650662386, "640": 0.0011521633386963458, "660": 0.0010345512272746527, "680": 0.0014703241871573396, "700": 0.0022689739261214977, "720": 0.002714861129921377, "740": 0.005800220039199319, "760": 0.0072751867583119864, "780": 0.012355766658887645, "800": 0.014635295468316744, "820": 0.011244771247342284, "840": 0.010248225188751379, "860": 0.013109683696949682, "880": 0.01479392021174634, "900": 0.016892151099455268, "920": 0.017663442143744875, "940": 0.01820524420826548, "960": 0.1109586885644798, "980": 0.10368549774809521}
@@ -21,10 +31,20 @@
                 "lambda": 6.0,
                 "etas": {"40": 3.772142176929246e-05, "80": 2.2706469625575782e-05, "120": 2.6326200988558663e-05, "160": 3.1264416616630894e-05, "200": 3.734036301684134e-05, "240": 4.846247119505167e-05, "280": 6.404192925339954e-05, "320": 8.471533678103274e-05, "360": 0.00011057770282582904, "400": 0.00015357450981873662, "440": 0.00020928137685536158, "480": 0.00028775575102761786, "520": 0.0003906344462048851, "560": 0.0006252858438352591, "600": 0.0008340419507301546, "640": 0.0011431077222103293, "680": 0.0017961608713683882, "720": 0.0031925221564149363, "760": 0.004980067265202975, "800": 0.008485425858496985, "840": 0.01032865422644551, "880": 0.01237070838947314, "920": 0.03340125151375621, "960": 0.08094067944719611}
             }
+        },
+        "l2": {
+            "T50_K3": {
+                "lambda": 2.5,
+                "etas": {"20": 2.250920776773909e-05, "40": 1.4891884818754058e-05, "60": 1.5223827915669208e-05, "80": 1.5639630260174234e-05, "100": 1.6268125506212327e-05, "120": 1.652049584452421e-05, "140": 1.7057065786420158e-05, "160": 1.715034420785078e-05, "180": 1.7656646576847373e-05, "200": 1.8257085026145736e-05, "220": 1.9303194915829555e-05, "240": 2.0251552070277165e-05, "260": 2.0904247090802156e-05, "280": 2.1963695428852718e-05, "300": 2.338154237461785e-05, "320": 2.518009516408181e-05, "340": 2.7632215436672177e-05, "360": 2.993943331246188e-05, "380": 3.2845862092756436e-05, "400": 3.667125617866249e-05, "420": 4.0582062652108116e-05, "440": 4.5668825691931606e-05, "460": 5.122729283710825e-05, "480": 5.5348235104926256e-05, "500": 6.332719154695279e-05, "520": 7.370348940305089e-05, "540": 8.64344362890925e-05, "560": 0.00010485919166295858, "580": 0.00012288155359968286, "600": 0.00014893493979579157, "620": 0.00016741987956066717, "640": 0.0001798707133915633, "660": 0.00020018856193477106, "680": 0.0002258201326020943, "700": 0.00027477537683857895, "720": 0.00031413701323857384, "740": 0.00035981607614058667, "760": 0.00042903544095452086, "780": 0.0005089691255931072, "800": 0.0005898799241695394, "820": 0.0006505113575950561, "840": 0.0007859076420017806, "860": 0.0010892936064263942, "880": 0.0013403461015629968, "900": 0.0018002949994522805, "920": 0.002621556670930086, "940": 0.004016384878178947, "960": 0.007658278420262655, "980": 0.015014497485306054}
+            },
+            "T25_K1": {
+                "lambda": 7,
+                "etas": {"40": 7.885411237921151e-05, "80": 1.695130914704497e-05, "120": 1.4020456470624625e-05, "160": 1.3330495163061e-05, "200": 1.3491282829324903e-05, "240": 1.4722835991976495e-05, "280": 1.6564887792584472e-05, "320": 1.855265423201437e-05, "360": 2.1481760991867215e-05, "400": 2.553426154865359e-05, "440": 3.027029618727462e-05, "480": 3.765916602025311e-05, "520": 5.0068265035909264e-05, "560": 6.023752262789431e-05, "600": 7.482492019220085e-05, "640": 9.622763599561349e-05, "680": 0.00013736714708269487, "720": 0.00016338166746091343, "760": 0.000247659357003169, "800": 0.00038514784378124904, "840": 0.000583681033053567, "880": 0.0006968971216342099, "920": 0.0014111171957347389, "960": 0.005987800101521067}
+            }
         }
     },
     "super_resolution": {
-        "kl" : {
+        "kl": {
             "T50_K3": {
                 "lambda": 3.0,
                 "etas": {"20": 0.000778915347623841, "40": 0.0007633780339799099, "60": 0.0007908141231719096, "80": 0.000788543595116803, "100": 0.0008288430294987044, "120": 0.0008392907314376583, "140": 0.0008925030169959898, "160": 0.0009190516306847781, "180": 0.000986800913723534, "200": 0.0010337262985591182, "220": 0.0011219889907389093, "240": 0.0011943036330360538, "260": 0.001310226177130967, "280": 0.0014183460616152988, "300": 0.0015733242833945388, "320": 0.001730196133012654, "340": 0.001937516604413998, "360": 0.002161490631568267, "380": 0.0024432350890119504, "400": 0.002708168867773674, "420": 0.00307449542457626, "440": 0.003515836769282733, "460": 0.004125363116570879, "480": 0.004754074570960041, "500": 0.005524380339583777, "520": 0.006439562376976526, "540": 0.007540773754592587, "560": 0.009500516111354843, "580": 0.011261532776188401, "600": 0.01348216019559666, "620": 0.01757517906334087, "640": 0.01977655744741228, "660": 0.024533713245498356, "680": 0.02746674570831025, "700": 0.03332074050640432, "720": 0.040243771626849834, "740": 0.04550001243601292, "760": 0.05033544385241358, "780": 0.061410491652647066, "800": 0.07221948341628477, "820": 0.0856409270444583, "840": 0.09925753198252922, "860": 0.12306635464984625, "880": 0.14476034450020075, "900": 0.17468273558090733, "920": 0.20411002595502606, "940": 0.24450559755448473, "960": 0.28248054410917073, "980": 0.3147428834175}
@@ -33,10 +53,20 @@
                 "lambda": 7.5,
                 "etas": {"40": 0.002727176113620816, "80": 0.002118729856897165, "120": 0.002106326729054324, "160": 0.002295490187260311, "200": 0.002577219155480812, "240": 0.0029055148578559465, "280": 0.0032434651627058597, "320": 0.003844541849902738, "360": 0.004323113740483335, "400": 0.00514218754706466, "440": 0.005974186955805667, "480": 0.0072513396445271225, "520": 0.009855938730607262, "560": 0.012369151964727335, "600": 0.015923742821581138, "640": 0.021501799745055887, "680": 0.02177964029551375, "720": 0.03612378754790826, "760": 0.04641127936774787, "800": 0.06494989916290186, "840": 0.07686161107101508, "880": 0.10411057668338695, "920": 0.21030777467686249, "960": 0.1347356123776581}
             }
+        },
+        "l2": {
+            "T50_K3": {
+                "lambda": 0.5,
+                "etas": {"20": 0.00028855906810001144, "40": 0.00018252251075077215, "60": 0.00014655976453229234, "80": 0.00012872307395854073, "100": 0.00011342178428976235, "120": 0.00010572723387176915, "140": 9.411218617644486e-05, "160": 8.871105518836539e-05, "180": 8.475913810054025e-05, "200": 8.131516973683073e-05, "220": 7.714109734750065e-05, "240": 7.620295840068597e-05, "260": 7.62307820527726e-05, "280": 7.303823321964653e-05, "300": 7.311525380813035e-05, "320": 7.95745068684325e-05, "340": 8.634892937583583e-05, "360": 9.627008476291762e-05, "380": 0.00010782826679609907, "400": 0.0001294239936281377, "420": 0.00014974205071155566, "440": 0.00016162172250863097, "460": 0.00017638081621832453, "480": 0.00020025535148243787, "500": 0.00022351568330417028, "520": 0.00025875363344605973, "540": 0.00032341707426654067, "560": 0.00042587316023395523, "580": 0.0004910637752685009, "600": 0.0005935274784633363, "620": 0.0007557815548553819, "640": 0.0009019441844744658, "660": 0.0010876850861333094, "680": 0.0012872008104253827, "700": 0.001528515131215205, "720": 0.0018386922484705586, "740": 0.002246468519145748, "760": 0.0027994054857132518, "780": 0.0037545329226670327, "800": 0.004453839350752291, "820": 0.005235482244357635, "840": 0.006663192577109613, "860": 0.00761341460916288, "880": 0.00955369356215832, "900": 0.011896386328399273, "920": 0.015882482051614927, "940": 0.021873072446695816, "960": 0.027160729169690586, "980": 0.13114472568136737}
+            },
+            "T25_K1": {
+                "lambda": 9,
+                "etas": {"40": 0.0001566684121864799, "80": 0.0001566684121864799, "120": 0.00013745984464489327, "160": 0.00010591145413933353, "200": 0.0001133765783443559, "240": 0.00013331100365410382, "280": 0.0001597699175899043, "320": 0.00021000760325224125, "360": 0.00028173776867778565, "400": 0.00035065806699962227, "440": 0.0004719081886972406, "480": 0.0007054373200230331, "520": 0.0009081616342901372, "560": 0.0014383193175793955, "600": 0.001935601097977074, "640": 0.0030684057501510283, "680": 0.004709319368864334, "720": 0.00786180984816619, "760": 0.013066580295808644, "800": 0.021678390591043335, "840": 0.03996850367727393, "880": 0.06557542936885012, "920": 0.1421624501010637, "960": 0.12574659563178672}
+            }
         }
     },
     "gaussian_blur": {
-        "kl" : {
+        "kl": {
             "T50_K3": {
                 "lambda": 2.5,
                 "etas": {"20": 0.001311277637655049, "40": 0.0011314792208201052, "60": 0.0013186921772428831, "80": 0.001161965484160248, "100": 0.0013730683554339122, "120": 0.0012312984890480417, "140": 0.0014760163095296494, "160": 0.0013462895586367368, "180": 0.0016383561831276197, "200": 0.0015191852141152255, "220": 0.0018704796563358576, "240": 0.0017668750655908613, "260": 0.0022029732696644934, "280": 0.002117440832022425, "300": 0.0026667582305060513, "320": 0.002615673065864779, "340": 0.0033227378642475855, "360": 0.0033246281530489056, "380": 0.004257323198960001, "400": 0.004339347273384217, "420": 0.005592718291095555, "440": 0.005808965844831082, "460": 0.007536144611100317, "480": 0.00795709164619601, "500": 0.010367841265905342, "520": 0.011111897162984073, "540": 0.014601192421942862, "560": 0.01572962570722594, "580": 0.020453046896395446, "600": 0.02245033051610993, "620": 0.028807566692203808, "640": 0.03205876293402484, "660": 0.041381670842137, "680": 0.04627552317583222, "700": 0.05792698423136997, "720": 0.0641696207414127, "740": 0.07855217157377192, "760": 0.08430869243734988, "780": 0.10214072576914401, "800": 0.11871801801180902, "820": 0.13689782336609055, "840": 0.15777806808719547, "860": 0.180580820528807, "880": 0.20387982661039958, "900": 0.2241532458656858, "920": 0.26195837225996266, "940": 0.30521510937946494, "960": 0.3967396717073996, "980": 0.40298601822973107}
@@ -45,6 +75,16 @@
                 "lambda": 10,
                 "etas": {"40": 0.00223172777962139, "80": 0.002324943896556253, "120": 0.002347626232887454, "160": 0.0025855993230245048, "200": 0.0027816403864307474, "240": 0.003232656561641102, "280": 0.0035120116186731034, "320": 0.0043515958429587035, "360": 0.005171759666352569, "400": 0.006783558909602239, "440": 0.008177996254973204, "480": 0.01063599121374264, "520": 0.014335847666818375, "560": 0.018337043722309928, "600": 0.024463545207151972, "640": 0.029903566958453323, "680": 0.043206551929219614, "720": 0.053882087120507724, "760": 0.06720757469247418, "800": 0.09491236004307764, "840": 0.08980582293523451, "880": 0.1322836370288231, "920": 0.3906382071142557, "960": 1.07926636526566}
             }
+        }, 
+        "l2": {
+            "T50_K3": {
+                "lambda": 4.0,
+                "etas": {"20": 2.8599508321212974e-05, "40": 2.8599508321212974e-05, "60": 2.8599508321212974e-05, "80": 3.232000934230189e-05, "100": 3.5889698442442026e-05, "120": 3.52073007132536e-05, "140": 3.906610368964284e-05, "160": 3.9779611206132506e-05, "180": 4.647805431999034e-05, "200": 4.5799158736207955e-05, "220": 5.625529605979526e-05, "240": 5.543457218295842e-05, "260": 6.974281948525453e-05, "280": 6.935233060522814e-05, "300": 8.90976293069901e-05, "320": 8.884480024377618e-05, "340": 0.00012015224040418777, "360": 0.00011764825438846306, "380": 0.00016261739904205135, "400": 0.00016162710355477537, "420": 0.00022402564101511123, "440": 0.00023834618237472715, "460": 0.00032319973499780956, "480": 0.0003303019556397128, "500": 0.00045877876121868267, "520": 0.00048307527073241795, "540": 0.0007146263571260429, "560": 0.0007446919190397388, "580": 0.0010589267468617488, "600": 0.001246586709515023, "620": 0.0015175041134437026, "640": 0.002029159451716644, "660": 0.0026273515313006697, "680": 0.003276171260052057, "700": 0.003977200695115054, "720": 0.005081355758370091, "740": 0.006431449862586164, "760": 0.009042150678260258, "780": 0.011377414049097362, "800": 0.015136768466397726, "820": 0.02109202217285341, "840": 0.028469604153117518, "860": 0.037300221385792455, "880": 0.05284632968022623, "900": 0.06871325698240326, "920": 0.10250500952531072, "940": 0.130896726703872, "960": 0.20719637640706742, "980": 0.23048109042050965}
+            },
+            "T25_K1": {
+                "lambda": 10,
+                "etas": {"20": 5.514891117948983e-05, "40": 5.514891117948983e-05, "80": 4.9385444334481226e-05, "120": 5.544310901077794e-05, "160": 5.871207283736799e-05, "200": 7.059294758318819e-05, "240": 8.414851116929105e-05, "280": 0.00010676823451650357, "320": 0.0001273334647988033, "360": 0.00018643723994848826, "400": 0.00024985589512634793, "440": 0.0003724955770220769, "480": 0.0005092840157396569, "520": 0.0007636186039346352, "560": 0.0010462538875478796, "600": 0.0015767989913558578, "640": 0.002745285900822035, "680": 0.004410929080168232, "720": 0.007625931463435605, "760": 0.011019866868347336, "800": 0.0230644453449094, "840": 0.03591833253921801, "880": 0.06917297408574581, "920": 0.13184045724004842, "960": 0.08889692816552142}
+            }
         }
     }
 }

inference.py

@@ -28,7 +28,7 @@ def load_image(path):
     # Resize if needed
     original_image = np.array(original_image.resize((256, 256), Image.BICUBIC))
     original_image = torch.from_numpy(original_image).unsqueeze(0).permute(0, 3, 1, 2)
-    return (original_image / 127.5 - 1.0).to(torch.float)
+    return (original_image / 127.5 - 1.0).to(torch.float)[:, :3]
     
 
 def load_yaml(file_path: str) -> dict:
@@ -110,7 +110,7 @@ if __name__ == '__main__':
         choices=['gradnorm', 'expected_gradnorm'],
         default='expected_gradnorm')
     parser.add_argument("--lambda-val", type=float,
-        default=None)
+        default=None, help="Constant to scale learning rate. Leave empty to use a heuristic best guess.")
     parser.add_argument("--output-dir", default=".", type=str)
     parser.add_argument("--loss", type=str,
         choices=['l2', 'kl', 'categorical_kl'], default='l2',

noise_configs/poisson_noise_config.yaml

@@ -1,2 +1,2 @@
 name: poisson
-rate: 0.05
+rate: 0.1