add demo code

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,7 @@
+.DS_Store
+
+.gradio/
+
+outputs/
+gradio_tmp/
+__pycache__/

README.md

@@ -1,12 +1,14 @@
 ---
-title: HD Painter
-emoji: 💻
-colorFrom: red
-colorTo: red
+title: HD-Painter
+emoji: 🧑‍🎨
+colorFrom: green
+colorTo: blue
 sdk: gradio
-sdk_version: 4.11.0
+sdk_version: 3.47.1
+python_version: 3.9
+suggested_hardware: a100-large
 app_file: app.py
 pinned: false
+pipeline_tag: hd-painter
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Paper: https://arxiv.org/abs/2312.14091

app.py

@@ -0,0 +1,350 @@
+import os
+from collections import OrderedDict
+
+import gradio as gr
+import shutil
+import uuid
+import torch
+from pathlib import Path
+from lib.utils.iimage import IImage
+from PIL import Image
+
+from lib import models
+from lib.methods import rasg, sd, sr
+from lib.utils import poisson_blend, image_from_url_text
+
+
+TMP_DIR = 'gradio_tmp'
+if Path(TMP_DIR).exists():
+    shutil.rmtree(TMP_DIR)
+Path(TMP_DIR).mkdir(exist_ok=True, parents=True)
+
+os.environ['GRADIO_TEMP_DIR'] = TMP_DIR
+
+on_huggingspace = os.environ.get("SPACE_AUTHOR_NAME") == "PAIR"
+
+negative_prompt_str = "text, bad anatomy, bad proportions, blurry, cropped, deformed, disfigured, duplicate, error, extra limbs, gross proportions, jpeg artifacts, long neck, low quality, lowres, malformed, morbid, mutated, mutilated, out of frame, ugly, worst quality"
+positive_prompt_str = "Full HD, 4K, high quality, high resolution"
+
+example_inputs = [
+    ['assets/examples/images/a40.jpg', 'medieval castle'],
+    ['assets/examples/images/a4.jpg', 'parrot'],
+    ['assets/examples/images/a65.jpg', 'hoodie'],
+    ['assets/examples/images/a54.jpg', 'salad'],
+    ['assets/examples/images/a51.jpg', 'space helmet'],
+    ['assets/examples/images/a46.jpg', 'teddy bear'],
+    ['assets/examples/images/a19.jpg', 'antique greek vase'],
+    ['assets/examples/images/a2.jpg', 'sunglasses'],
+]
+thumbnails = [
+    'https://lh3.googleusercontent.com/fife/AK0iWDxaRlJeZGIBuB3_oGKKhd5buKaL3kJ6moPp7r6svDYFkehrv5XKyF6mj_Pqy3yV-qDQQZj_n8CMpuYH_iDy5717rPL-qpXf-prcIv2pET4LDjFInFQVLoxuurB3_7fugCUogt5ZYIGlTgSbirJkHDqN5min3riiRJLd0ZuGN-ETDDCs5e0wohdX_Wl_Kv5RAdYjZqWFGfKcmzCF-1ny6bjCab-1hcDzIaokggkl3INTG23nhSLWhNB8EdeCdfkQmoF5fROfCPe0Lsvk6RwlAr-ZQ9jaszJ355oXOz4Y-IQLfWvnyfdyxQ02anJF7DiBZbfhH4WcA6faK7Pbjo7RIt-nTN6THwTGlEmxHEzrO-1iuy3j4UfMxPB7r-RjDrn7M9KwLuWSybIJ3dgSx_DF9OqIYkHeG_TSvs1Vi-ugbq9E0K20HNXkFlhe6ty7Ee5xr0nNqhD6lVr-6wbBvI1SJUQg03KoJkduYbTTQ-ibGIwCu7J24kdlo-_d1xBWC01zeW2MfqPjpdNHUtnPAF2IuldEsWMJMHEpWjmXwYfM1D1BmIcOuRLvdNEA6IyPD6VCgLVD27MdLWxdKHgpSvTZ1beTcMM_QuHV2vMJnBT4H7faPmLBe9yMTKlIe7Yf8fGiehOrfEgXZwNyRzmbOyGKfNRKiSVuqeqJt3h1Vze0UKFzAx3rYnibzc--58atwasvdKY_dj8f0_QKR0l2vzWogqh3NdhO9m2r77ni5JXaik25QQzF-BPe70ikVgEomHa4xySlf-Gr3Z1v_HWDa4kYg2KE3P0WjqUD_cmdylbZ45TIOkGAU1qmiEcTCs_wOkOIfCC58Z9P6Lff_BRxc_lhut8hp-Pe8tfRXhITYFFRthXforXyDuqzPmWBBz2EnUHqxa1aYOo4WeQc7KTXK2kF36qAzPwm2QFDreqV3QmS02Gev8MCz2U3bQQa9H8VSB4ZGhrzNWIwG1R8YU_G1Xb4BAgjYEZ5qX6WJaQrjuD6_Zw_pDxRew8t0mjj_tjCrmoZjpxjHsgtudH4IBah5xag0bGdUSThnszYJPM5g1weldimKE63HqaQTG-IN_N51nBken4K_0-liw73ABzUiA6EJzqCQKEoT2pejNTN88N9RFXXB5ZJ9x0NvtuMcy_JsrsVArfA5b7m4OGbwF6b5wN3Ag3XNQ3d58hVJ_Hw99HNIvrjTVCVmU0-DsYIu_njIfyjdqps1cyv6_f23F0X-q4ZsbooPoNg2lc3oeFtC2K58Dgr8JsBjt33Lnmra3YG2nBh3lkIycCvrUOS69xo8Aq9z2ODklCd9soUUQVNa2XKKsMofRi8ESGwuiWYKWdSI7XAXi4dbzFQhWwMRfsSAk3KGMpnUnlOd8Jx68fiGMwTKCCsgIPgo8mFZHSaN29ipzNoACG1bd0ZFX2qxa7UfMnnRobl-AcFLzTtOYD2T2PcIKKjnxy1-gG1Ff0mgGS-BCDl9TtCscAoKcXqTYWCSY5otnfVPcjgEFw_0KEmIbmrf6Rpyr7YGJFPTQfDRah5Ro5LhBIbXhFAkyyvWQZCecGWi3lRAh-pSm2zuGl0nVdykgMqzVwiat8_lhnHMBNRg9xWyLWrgw3bWpdf64Bjgshr0V1XV7kikgRDpyPMJPDNzX1jVASiL5S3dmWDnd49tdnaKBUXGAIWmTyUWfs1bDFZZoOfrmLrnLLUT9R11lMO0EgkZ48_UU7CMJKwgq6hegh-ErsV6S6_SLa2tQYqYlSNUESO7jw1w4_T9KTkM2635QlPH-A81yvgrSsRo0lq3uPtFEHP35-dfGT63yFd=w3385-h1364',
+    'https://lh3.googleusercontent.com/fife/AK0iWDw7CKs6uhv5cJnJ48zW_5UW3IjrXwWqvemhbJGzB_xNhHgNyt4NIP0xPkCK_nyGFEsbb-4iaeWnoJbbeQxU4yBnnuckhD5d-t1qgOuFke8NRVaGCDgtDimbwNY7Jkkc9pj3pMd33P7UhBg7rVueTL7hmAyhtT_wLD9B5a50VfM_T9ErEeYUgsiEYPoE-msuALr2Sl5t0jXttYNt8R2JHdXG8Y6EegiQaEdZjZbJxdwA5S175pCmrOvVVBH334NwarLF_HDr8ESIVEkYxNtLDwb0QPDZSDeci-JS2WZcvEQUmClbxuk3hjtOR0Rm6VX_wXo0-nczybNDdTiZFsBc1Fj1L_08VLYP1OIz1fbcRF7iv2luyGmgUijGaJLr0wtW_C65eXK9aU-NJ5KVt67JT-TncumCqBpvn1-msbzbcL8krOmYqiZthOcvnrIQHeW1jIUVD_zMZADmuRdpjxO-Dn3yC3uev0ve0u4b3vEHR4iiLx-4Jf5DrMvpfdHHIL76verfqhLkiz3gtZENg1jsRTcjH35AWspo3-lMJnZqMP9wzvw7ubBBP1QrblSwfflz4aAIzmh-WisQ5UWMSltQ9DwpT0UggiXe0JbbtloWkbo2_VpaaMJh7VhBQbvRbFUxCm__UPVfglTfdiNv1m2777oGwyDbv682I_qGDW5nG91D22pqo9-enRGLvCnN-STKtWTnJQ5Qod2QYRAEI-1IR9_h-UWtCyBLLpqcGKxkHaLZvjpDmTdiPhVkoG1irracCLbPJGvrrclorr1k0nqTIhvVVH8pNZdx-yCK6KFFGNAz3aSsmxGJURWEt4TQEVDLfet8iFzuLfD2Tg9gJ4kozB88G8PrfLnGopwlPO4y6GkcfJggtshQr7yuo3xnxtci1FlLhOJNkMKCg0xhL-tVDIfWsMCzSIk08XtGgitU7DK8CE7DlkA7chCuq2BHUTeFCEALF-5DLFxa-3ru30gtNsbd71sNH01S8qhTVegWM8yQwwXXhIVqoWM-e-SSUyQldgklvatev1V6iDxLh9u2nbhoxxCcVZvyKx7rUyhSEEQzXF6nYNQn5HaAmp4jW1rCoVVvVMNP_wfvY76vamAPkw22z_nylCpbKW0BQmTBZStyoB4u_lZANVkL0fbyWoxyfZHKOoDgyTKdzDWmbeiIPmr7nvF8QWYsq160pKzIKIhj9ccF2NqyXHqjaDzmqVfZ6EUg0sJ8CeiCUskbSnyxZHVeM5GgG7cj9xMncUUsdvd0WTw3g0aqstb2MAxF3ucO0EXP8dAWqU2XrJ6E0rA51Jdn16CW9KV5HMGEcB5Y6bN-Md8-FT1r_0NQq8ZOJ5nmTzg5JNSTEw-FOBwSvukEMDmbf6ADs1emocozcqk9KiYv1ii3niD58XPIhQWfBcgyYBsqbJQ5x-UqijqcihN7i3RgQMOVJem954MKU8D-DzSet2FUcGWseyzF6Sr4GJQn5g0rOXtFP2HwT2kTXy_pqad1ukQkovfQG30gbRvIRAzMmigVJskadvadasF7Lc5eo5Pm7CPPZQ-ZJUlZuowbagWm0Kz_T-PIreVNJ5WxxQ4w_HH27QqKxe_LZvHJ9y74O8oVCywVJeKQirxDc-yKuUNJ7vagJ5a-paB_cXf6RECpu7caFM6U3g43xVMtxDf-3d0r6G7zs0oPHb4uDbovgLzqooAnjNY3_RGWldXUks0pmCqFbEhLgl8JJd4Lyv5mt34WTDqPggRnraAV5AEtZ8NvGNDWk-ItblCZQ-neXYVj0z1p1rWSJiG1XaDa8ro7G9-7XJtwfJDnFZsf=w3385-h1364',
+    'https://lh3.googleusercontent.com/fife/AK0iWDxK7tdzr1tx6g49uS-pyR-7BdPJfk0_wWoErGJiQPMRbggXquguwth2I35go_GFsW8TUoj17p-jtCKCi0ryGH2gTGNvZEki8xfPn1aroIxOXRt7Ucl64Lu2cyixNSBqyJUBQEOL6LzY6DlwXJ5SVxRRvJkj1Vz-yJYz4mNA12YdDFPj7UkLs_7PEDRanydfEVFhIoiMrYitDW8fGsqn53PHaB8GoxoQHAXFyf82HPfhqqgKNRUR9yz1LAN5q72ERzxg5h4apStr16aLzijuuJF7wmtltf3HdTIMwAAd8Blsgk0d88rNZzItVZMdCeHgAyvoD7JKUCTWWJbu-uXf1PRd6mIg11OzweG04c-uloZFvJdF0pdfcfocibniKt-filYASFN43KYtO7Eyzc-YUHJn1qm73eDr0RkntHy1kQaDPmlRvlZrtbDgZIEau_FsMa-BRNq_jHlTDkyRR_8dyBf__na0I45lcChJWZCOfinbQwzQfiryK13yBB1bQDGk47JM5PTIEBK2JvsOEwezA9geZWhi3oIM6N4b1YXzvWbDYbN-cQvrBd6cuXoecUL-i_Qqda2-xGByMP4BUI0J1qF3mKwzTLvzZAk1f-zCUNHpOrzF2-WsXCUoL4R_t7MOZ2OeT7a5pCAnyiFX8Vdq3x9QzrJjjTqvttS5ElXVUGi5NBVeRV0hbXI5XLMKmiu-lWV3625eosYK0FE6hmFZk9ZAeSkrLDjkh9k1XwpB2x83kcpC_tkk25G9czy6utoNhC2YHc1uVvnUjZ0DjK2d0naLVOCKOCVRcTSvN3kt7PaLASArW5vOzZX0_LcjuhHSzK63mok6XvXXZ5AOiF2fDhUNUl1W8TaDz7aIssbibuLXxjpluone6WJglRrxAkpwtIw2rwk73icupAeVG6Frx2QctrHC0vLG4PKSfrIrZRiJC52Oxpp2dAMnhj76CIwX9wwbj86uSyIuxeIviiKumwTolZjbOKhPrHY-ZOydVDn_ZUsln4mOfDZXwUl9p1CanzpLCUsZTXcq12zTbHJRPwNw6SH9srQc3cTYYsWpBu77VmS5zVkIndUPItXWmUqds_2AI8LCUSWE9NVHiCRSw-B8J8j3SkqOD95-np2cMxDxoVX4nD11CwBr2W0xWS1kqc4mZL4aHdwhKdDcSKk6EGG4kmY1Eq1RsYxc0I08TJK-_nrbWTgA4NDTjh-oJpFFiF11ZHbEksKlSWBhH-MXF-0zxmiar-EIhAQe4hQX_suDty4GXxMwzcF84JthDAvWC0tGLJV5WQGLkLvBkTODythSNoN7l5HjnzoJf2JKI5AF8W3HJJWgntD0F2iFe5K0Ik_huovMbzGjkgPZDvWzrYpA2V6VZyPI2q08axvAfaevPrCd7G8Jy_gliK3hj3qxjIvqJXBUSO2puqum-TlUaSYgbjhWUCLXKE1BH8RRDt0brLklVvNwGpHxG6Eg9vmzp00fV3qhjoJqokTYOAxcAumtfUpyJFVk3cvpgm4826o4kKUM6eXRz8ke1L34YI9Rtf7Ppk2864hIBIy5xA6ajIpI_rIXWG8ogrhEp9GMbXGkLGiLMfd4t7P50JynPc9hYWfachaTlQfQWuz1NtV3ZGdTaEe5y_SpEwWG-UQ-_MglvY8p0EYNjGc6yhu4oEx9D0RN5H4QA2e0YjmVybG05Hfm7LwEBMBvHb8GHnpHGTxL-WDlZtU8DqDmGmhv5l4npcLnuVmSE4JcEtRJ651ZYn88wD2ghTAJtcEuyFW3nzZ2Na3z0LW4Y8Y42YdWb3hmgGsJAYBIa9YC_Cmh=w3385-h1364',
+    'https://lh3.googleusercontent.com/fife/AK0iWDwPG9L-g7hX4NKVZ45PuqnRkQUV3XZlMEWZxXASV2zYVytOhbhds--yBA0ZXUzxeEAbpGa3qCl8sXlu-of8ZjVVNSdSen-zgpoN0BT-R7JGmTqjT9aEdhnysy85Gqr4e4A5KPrInLLvCCkFSHolBEhu-hO1u8kEZ7aDuk8FsPvchP1SvqjnuWtY_OCYa2FpjHH_i24cITjs58nlDxNTFdQQNCnX4KLJQAzh5cKxtG_7pqoBfBzpXBvVDwoxPafbFn0X0u8oJ7V-VOh7faO0JtJrsfcdS06tvw_J29RbDdWojpUS_pRUtA8w9Z6flNlbShj4Ib8X1V6veQLrEYWi3SojO5tfntpl5bKXPGgWhsT8mykvfhg4Tq-Ti9kyPBDPQqEqf0ll9-wFHYoAWSMsmCvZ7kakuDMH1766rOk8QgqYbDr6kMuJw_OFBR2qX7DaQw9XnJqGv7J3guCj-yU5vXes4sOtZ3n1IOheXnlvJL79KhXbsLgznYBjC9uv0fDLuqKaFL34YisLY5xY_2zi_uzTc5BXmtoAFH8otrUVUTVyt6sEPaCtjyPzG2uoSq094QFY_FxagW0E6OYPqskBUwPzg_Wc2eBwazaGy4MXEoDzgvK3PId5N4MWqU232uHsUrEEaCUUKm8-KX43c0C6O2daqjwsh-bxKOIic3pHqDoAOcq-QA83qB7pPyWwGddsaOWRIdzf-QLrB55YuvTeTmOEL84m2YDtxNEHWdYcnXYlZEXAex1xMOfqkbCQGM8jSgC2di_794HKMpsNtKwYH31WoI-Pl73t4THuq9CX1pWYdjhH0ss1j3PUMJ4hELE187E2m6fhQLHNNSHRajfesIwPVP1FgP0W5o-AKLaC7o53R0XrOZ7SqYjOua5TX0RXyJE6ceZATyiZ_2tzUpc4baqRGJb88vr8dhfECh_1J7O_8ufMvYL47HWhVNqltcIGjujtIXp64XEqShlut8TkqtB1-3OqE7gNwfiP4859pPkk0q-kfIdaLCjqB6PeqNFNgcX0dfK_-weXpM-LjrgNM_aVlQEBKwTDjJteUHY3pFTrEU8FoPnPeIjbU_rGIcwA3Lf_NE0CJUhKG1gzKYWG00CCg0srvQRipJpVrW-4SkUrxtW28iq6kVRMER83sN2RnVi7ugEuZ3S-OPdgXkUglk3bIz9ehDcFfQrll0iCabIGQxrN9-7GfKSpi3j2nU0aLeZ6DxgQ7x9f9f-hUgxM97i_90SX-S_M6Lo6bA28RB515HqqXc8FN72Wsp_XlFp3oTb4cJtOI5SSUYkdxGtrn0AQWXuE0Rp1DEaUDEdVbu757FQsct4Us0jLByauEfiZcQS_lGTjzGQxud-2NMKIeIYOWAxBGz08eTWv4b2k_IsnekNnXAHzP5WCYFoNtDiGbgURj7QWz6Qh8HbTsFJCrI3mPYkTN2jMUwviTUKpeElUwqDeA_DhT5tiHa7ldtdzncNzpxH-J75asams5J_O2W5dJMN4PYUxGmVw5mhWEClFo2stMJOfPkrmaga4D1OXc_C3Utf6OWB5CBBHGNjfAekE3QWm7ibEtwC91g1pIujyCUEYVs8YiFi0RWcMsmhWG2yrghA9Hu3kERWuVT0nHHfLRx8L0_PjlQBkNijUjK_gI-C66729qLNsmdwZxym1JCFgV4xRT7Vu9EQGL3tyhbOLRWYlHAjBE6itM-DM5T7idIyWamFBb9Nt6ZFehCpslKzEHy2VtEyiRlZ2z9kH-IEIuZ2qCu3kiGL8m6yQTyboTCI7LNLYbSn497VR8h3WkcqWqlVlWOzGUikL=w3385-h1364',
+    'https://lh3.googleusercontent.com/fife/AK0iWDyoheia4yOtefRVFWnA5EHolnS-xa4pPy0yL4Wb3uIr1-mSHxQ54i0wKr3lk6uhAm5qjJ9UVqwut3UqR54log5XVnaeNu_3y9Gsn44quU_HsGAY0-84HygWkr9Ld0_Dt1JEefD_f82Ijp1TXQf1CS6IbruUOnrljFOraQ2Bu-1To3p2Pk-T3tU8xVCxU3pr6zvFz9UYHFTss8_Xw70ZtLRMT4x4suHtOaSPI42VTq_T8HEnm04Ie_0Yh0Ri2_P-qsaP2ysz-Wnw4Ykbj9nc2VIDqtvCRwti36mlNheyg_8xOLD9sMNWDu3PXoRtn7aBUpw50GMCqeGUMAUMPKhJuZoTDdQK2SVHJ8QwNhYcC8mhAmRFvt85hvuT6NyZ00SeYwyj9_rux77vZThx5ioDoUAH3CQBcwgH82xahatReym7ehL3DXm3JLHDdQLbRM6xvsE0X0MsMXkuNlx8wn6HyteW_yq8fK3wQJ_3XLh-gK5YOFdvd08A7IIK6qi__-o8nvEK_hfHgMMS-O9eT9acfa2Sr0rGGNvoUlpljVOyONyQftP2nGD92R51K2Xcq9oV2sjSu8TDDel2t2DYr5gMB3FsDfgEhQEWE-O9fRLkzIZnOTTAcUDoS-b3R_kB485Ry56FzKFbz3w3tvHhwJ3W-sqvygb8LDoF3qjURWrf7Pau6UMjTSPH6FTjbVzZeWITKsSnA14xA2wj6xi9Bp3JkCOsT0qOfXrPkK7uT3H2U1M8uqFjpNTj6u3tZyF8GgueprmH13rJjYst9d_vevhpXpXIhSDvAbJuA3xG-YNr_SG8BMXpi7N3NC0VHmBXhl_wDBVUnAD6VVqqNtXzB-6NdZzjZKnxApDdi5SGp8C9kDd6bkaXUmwG__BRNVqdbchMw3H1re5t9VxiqWTelfGl6UqAX1W8RzQR21bgu1x7EAGbVsC_UpMxeDeJq9PprMF9cCRC9ziT_H2-ubctN7O9qPpADPT0nqWN28vH-9CqB4jPeBqYwi658twIpLwRFvEukajsrvb-OqmeesnT9QpCPEpL4G0HrjB9rkRX7g-T6q3kqbGfvnqgW4Q8ilOUnkEsFK3qLCIVxwp6B3t2yg8XSOsM3tnzsA1Ua8MofFKqvmwaq7QbIBcMOHa50I1fxRMs4YEVLgu89fSZxRrKSr_8oUXcRWiqSgF-pLLU37GYMrn3yXUJxxO4bXiEifeK8id5H1khO-8ZEBXzwuZBQBYLXbCkCou6enZu98tfPA-prr_NpKfSZM4e0clZWhjo1761-BmJSlIG0JrTo2N2cKhVz-WM5BZjVr1FPYOri5fIjORUL17_RbqMw5MefYN6tLPnpVrOSUmKW9bVgdFdOVpj9Wg5lZxxswAM5qK-wOzEjfdBCW0xjKzxD97zhszCKxc7Rj2uoaJzk9CeauU83LYcihkyMHn5IKhLeAou2yKEwgqXkU0LUObdUxqjavnVgVMcVYRds_j7zpmpKNT_KOV9s2jus8aptJl8sXZ_Gzw0vi7wC_AuAGfmNCsZBEFhn_b_ZgONqdaR9EKwP0hVRDNw4dZkYn9MGqeiX77I40eEbnwVbKaUZvK2Nrt7ukjkgSP6FdvZFfs-aVIUFMc6rBAknAFDFHPzFYcy9ANPDgVAlms8fO5GGuid8kgpxtjSoneUG_A9Az6JY-suY8zFr6mDtJC1LuY9ftIKW37ZaMtutqhbBX5b5w4DLnO32Uv_ZtK1nbV8E4T0KY6qHm2kZiIYuYCsHysbvNW4MODUihrpDh1UfILdMEsae6zGpCTQ7gnKlrB7QJ5Ig28y96Om=w3385-h1364',
+    'https://lh3.googleusercontent.com/fife/AK0iWDxriikbWcV-sJ5xBcy8xcJHsRC9EEBmCimXzsrBiSOy-UnRwSoGhdXNqnB92vMZk1LiOn0h3KbBMgbFr0I0SmQG5TwG__bM8AvtMrXA-DGAaNTktO1JcSPb-wVr_OepLK6P1hHyGYSvcGDdF03pIFNcKxN6QLCZB6rgFLdaWd72z3Dx8eB8QLtU8P_4G4sT2oJ0hAUlz9mKz9lTxWDWl0-1ufKUctmvWfU8EjNuQQzckJA0YwvV2jDl4ZA1r42UHDss6dy6hkjmPoDZN-p2UW_3Ju4vOVAdv0Pf73demNu-L1LALuK0rq4d7OHUaAuP0bubXJAH-wsuVervwPQDEsmBwR-FdW8jfdppKxy4MC2ISf-eyLmsTYR9dLPIlKkOAHh_84vdLeWdGtxs9gES-jhrqOiW-brFtIZoKbH_nR1yLeq9IJ7Z7-GJk3PVi_Ex7gT9WJyIuySNi6s7GH6AnDFf9wfHkzyJ1qLDKMddrNi4GfEyO89S0yScxZFW7hERAH7T2_1YqkeMv48ik9dMA0RcpJK0LYA8GDD5_MQaycXUjeoCO5tvlGQUEE8Dt815Ev3xh9dlKJiKJf_cClK_kL7iBICasjcVSNxP606Zn1fBTc_hF3QymDP8Q_Xl-g9p-ufobo_x2x6nOOdfiq1q37ik_3kPZCPPCnI2OtM3EYZ_yVlFmwWbtqxp2Rz3jaWo6t1TL8PWkFZno-aqJ8YEm4ppZVyX7ne6GTFKzFpW1SKnAnjqu67LjS0DCFhFATAcYhWRmdb5cMXve72eU5DNLgKZOaM0THDa5dnQPaRmu-c-7HlB3WISebcJHj0vIDw7DLwxaCnqqLgybvqW7O-Rt82bat5Lc-jIyMjjkZOutnc3OoYTPRN5PrLZ6HXGFBTq_s5fCimxpXvlw0bzNHqOzovgP4NC6UChXwn9CxSrbLou8vuqeD4YqyjBhh5Do3l7KcGMZtYUUMUhVf7fUrZIJcLb-ZBrg6UdiPRc3h9jcubLjXcPIrzxZeqgQ2ccJRljZqq4CBoX8WU6AiEe51Hbp0C96693G7rVomhzxa8JCMCCL3sy8v8nl1tqfqQ3kE53XnvzuqNMzJCNIfZg-GehqMBmJZp-Vaup9DOLWYWnzRqYsO9pC9r37Ajhh-wkxZyV0XhXD6EecYGcXBX8TOQRF9OMV9BYRDYVnpnGfTBGyxI7lkCU3prZQL0H4JlS48k38oWj9fE1F-PnGysDXLAxYqEERP-AXIwgsPyoKYrq7mG5UzBcEn6qminlmT2wM-EnInAb7F9e2aMxFwWv17RG2fWkO28p4t9hiMEyvAQfYyWKhGui1yqPnHnytNe3BJspV0uekBUBauvvHHLF0_tgHTJwcU03KlOLto7iIOUuwLgBT3z-_diWu_w9WTiOqZppBYKHUVfQlR1JsFn0j0Tg-kk5NbRfRU5BB4nMtQlkpW3vjTrQW88SrSKNOr_vepd1F39EqxBNG4tIyR6lvSCSGefWzUxwHBoHou8MtGPvcdxInB7imHdpwCFtoVy0wGh2kXy36CyxXqV82VehVvmfc8AnNtgVWTPJ8vvs2AQ2WL7xd7lE-DCfy-la9d31t6iLxVAkq_S067Z3iyTiJg508a_BGt_UF41VW2m2dW2v8KRjgXks4Wz7kAvO7rLLhL2CMdh-W7bBP4bfVa9BfcKVKGxaa6kCQB4rXtFs2gEbWx4o6iXf6KIUgl-RAdAgYYjIsavlkyihtgQ28-1--7JjaeO61d3wLYgJ5_t_POL8OjmFJfPeO1m9HTzDJf30C5hVP348dRhrkxSTPlSo=w3385-h1364',
+    'https://lh3.googleusercontent.com/fife/AK0iWDzcrCgKc1IMiriL10KOqvN79VCXca9U5g657RP5HU1zsUNVcZPyefBpRCtevRE0k6FkGxUa4yKW5ELtqaSmVRE8R2jwZPHmpd0xeBvIHYeoySmntJD3wJl_iC9Ma20qubbH3OFIxNFFLCXMJDPY9cJ2D1xFzjD5t6jQi7tJKXc0B5W9vGTyQ3dzHlZuLEfTp4D58WszOFBLPsOq8zeve0ej-2wCEMkCrT8kTwfKTnsi_GXRpK1xFRZjczbA9kSwIxZ_x1iWNWW_XY8aw5Kgn5Du-4r3rBqpHr4_fzv5ehY20T7Cq6Rf15zCBnBad1HGxLLnrJrXujKwxHELaRfqXPItfQAmuoIfL-tgK6giSSNKDDT2Ynn7AXxHz5vzD3m4ALYWql9Qpe1Jw69AxGZt3vyuSC76LPhMmNIJzVIXVMLnw55hHCLH28GVg2WxLTLyUXoX7o-PHhFjGoBbx3X58yQjyrEQPbfmY3gtuijG_vrjoz3SPcoG4eXB6d_NnEM3b84Ml9eqE8BjlACMEJriwtqLEstPBhIsJygDZKB822bY63pG2hYzYi5-bWS19NRwun-jdNraiye8D9AIqxrXaknLvbCnWbgbKDJzlRnt9Nz8GJKU__oZ1wQgzq2DOQqEoJxPwpocJxwYjnzIb-dFqhcLkR9SQVx7rAEyJ_u23WBSp4AVTw6c7sYxd597C5EFBvzOv4qYtJp1G-hBrvHYFvYD13oDtgcFDO_F0nXg7XwmCJ7aQ28Hp79dWewQiSq0nBcgWSOic1Q4feWLFEL21Dw0pFmSVF9f3C67YNA7ZXukOAnv6xUN5RNVoBl58-KDOFSQZYysT8FGYaaY8bFRvYvI-VXcBtMuAw1VYMKLKHrhg2mSvvdPCxsGH_rQWijgGUm9pvqbmyrPzYRd-Tr_i6pKH7EMcBYgcvZNMk9_B2EMj6NUWLVkHEUTDiovgFQJDrYAuiUcKZgp_MXJcTDO79qJEsiP52C83C8Vg91w4t_Q0dPklY8EjjpbMZGvsr5NC-TlQraNFszsrUA_JsNdnMW6Te2GMnXHDaHjEDCdX2kY3XC1Ltzvl3f4tHY12OMOUzMKHFWHWyBNYNKlHnwPdxYNxvHKa-9p_okvxw547oBfsXwUrRpQQVxljmLZJbbGpxkfbEW6Rg70MRKHdbEUg0h0FS7F90kuD7pR1zUFv92fPoI5BUNUn2XjQb-DZZaC1oF5VtRMgO8RSFM_Dolrx_c6ZrPLy7bllakKTDj48CNvafL3UxaD6x9FFu4dNQVHKidcVS__EZ1SYMSZoDkD2sna5Pzpjl0Qz61U4KVAz5lSCJbF1stdiwO4jMwzmAgMNV5-fnJA3kkf9dDaIzqk2diKUh-WGSOwEwwizH54Y-e2EASTqCzGK7FGoVRVr1d1PdN7wcd8MuRMXqBfonrIGmf5cSiuvOL6odbSChO2WFOKkRbBHeV1pO46uaBVeLKGjcbAALymYsv23_veIW-RQdHnaVlvB2HE7wD1afI-3LaRKCjCUD_X5QJuTXff_EQayDwtW46_0hYRI76LGXjQuOc_LUdIi2QHlO4kt2kR8eLBPm09gKl0GJrZoS9HxKS0sERHCjUvugFcvVmgt7idjFKfi0AuUz9XmqYCoSRiMYj9z612Ot5L2D00SLTOcTLV6nlx-PeBpAyiu2ia0ehvVjzLUn9ai0_XOg1bU-ab8fvzrRJIrao5ZxdRp9wF41lJpkanwvurynNJXXu0uk2SoA1_soIeLshOsONxT4DG_PitVFMYjxY8rgx7hfNBJmAJW4GZPsnx1_P5Ojip=w3385-h1364',
+    'https://lh3.googleusercontent.com/fife/AK0iWDziEIX1beA1lUdgcqMcnVcmRuODfH0IHpIkZW5YYZzhQRcRmYCfi9N7-vRWfcfMTuR28ZWDA5EngVjUpwIbVxRBF2DaIb_lJhd9zG3arGxRbi7CwmWdhAXeODEvniYR-IxtWB9lYpNd9hJ8wdleTP_ai10Xscy7iFeXOmFZ76dnr0r690LnZULOd7iyGv9EWZmhKhw4wtEJsqi3e1Yu4CXRsVrM5KLYKG-EWdRW_-m9H3o0G2W7KvOwDvqAewIz5zApPBHvuEE6x-XaUOuF_FuYQVhsKpfI_1Y70SCjOCphpbn51Bv8idg0tgTDn8oL8hkvSl50VqgQqjCLNxmCHlQE88xmjK_4NMI3kbIBLWfiPGCURr95dZt8eniqVi7yu8LNgkaMixAdRBCrQF_z56EsIvkozenXBdS2FiaUEh7LqIHLBcOa6ZaxV6_3t5Q83wgJTaM4cNkvH5_nCeQ9wkwKjf7zcBxFusa5LvhM-qSm4BJz3WzE1zgqTLVnDeh-EFNPMilPhevOdBuNfTY_VF8chvWS5Nwwcxlls8xSdVVqblYGw8YBlzWdi_X5PqynTKn6aWE0IiWOzA_O0hg2q1FtAHRT-PaINo4wjIbBar6fiNNwcZIeTKJrijHcpkIhnI8PHxrUtO3s0c2pfLFvuxCCRMSEfxpcwt0rz-ODEWIZkALajwE2SyFV6Qioc4fH_xWnI-jgRvzHjDf2c14vx3bXjM_gy-25mrECLQYcSWWZVINUbvKf6_YQDwwzAKL9zhMpyGa4EToTBhMSmroGi-NwIPxh8gAfdBCAh8TFAdg2aA7D3N_KpAv4Eh5bkovhCiALFYkGLch6KogZcn7NU3OX8qyn_wJ3oGO2CFmfKkMtLHqmjQpnLtM1U9BPnRELir7pNyG9bTNzs-Vz7-Hzu3vJavGeRhypl0JCoGfO08be-ee_7EnUcKSdepfd3dG39Gc1eulLgIVCRb82Ga5mAkNp_SDSa2BGI24--uOyAUwTBazpQjJ25W0wsHpLRF4obk8Tygl8Fgt2F-VPXYz1-q0x3_KZVWf-PJmKjYD6t3ICuBMoFeJtQUxp88WlSKC7KvhEZYdEaHmEabNNK7j-VTAgi0BeBaw_dTTO0tad9rXbCW9Co3Tc1YXv53oz96VURj-FAKHk_PKPRSV7-NO-BWAk1DOTq3ZDnlKUTA5-x6k4IR5HyNzW9C7rIPGzd_PRA9ddSiRxOjSiBru_P8xS0zQn6p75V48ZkoNsLPWEWCKhANJOaOB7Y01pg3wjjnftuxkp0KpokrlCZVUn2eKPmB0Oee6TP_6DVFhgM6ksqLHO-sNxpehUjWDx84znkN0MihGRgl6TK-6xnWzD9tjvIOsK0mBzk_XY3Vuvb5OEZvLzDJ5POqNHjLcAFaDtX7gsAUtEWk20qmRbpGBnHiZv2kLOUWCy6ICkc3yFv5uUMx7pxgfc_YO95ybO8-FTDG7m1yaoz-WdLV3tHao4_MfFaRXGKtV0_7xnlyXEZ3tMYwKu4hRx2lIOsL4Aff_O8-H0jmJId0llt__iOdVDkuypQWQDOKGGP9B1_gfLkV-ymEP0Bl59jQWNnAqE-jUpTeRRcUB6FkcH8XBPKL7F9N0sq-6XeOmPPpsecmm3SflF6zJ1YV8Uv6H4_9_uQLVBB8wXSvtcQuwgzYnrtpjpMQwFqSvJDhcCPGRfRCR6H7oa-T_ACYAMcICpl8felwVUOQs4O03ywLHNrZBY05hS13cj-_aYw69kw9TdetT-GbvTKC6eY5uwBTq4ytb4eeJQJc4zBlB2Dw1vKmcgIFfZ=w3385-h1364',
+]
+
+example_previews = [
+    [thumbnails[0], 'Prompt: medieval castle'],
+    [thumbnails[1], 'Prompt: parrot'],
+    [thumbnails[2], 'Prompt: hoodie'],
+    [thumbnails[3], 'Prompt: salad'],
+    [thumbnails[4], 'Prompt: space helmet'],
+    [thumbnails[5], 'Prompt: laptop'],
+    [thumbnails[6], 'Prompt: antique greek vase'],
+    [thumbnails[7], 'Prompt: sunglasses'],
+]
+
+# Load models
+inpainting_models = OrderedDict([
+    ("Dreamshaper Inpainting V8", models.ds_inp.load_model()),
+    ("Stable-Inpainting 2.0", models.sd2_inp.load_model()),
+    ("Stable-Inpainting 1.5", models.sd15_inp.load_model())
+])
+sr_model = models.sd2_sr.load_model()
+sam_predictor = models.sam.load_model()
+
+inp_model = None
+cached_inp_model_name  = ''
+
+def remove_cached_inpainting_model():
+    global inp_model
+    global cached_inp_model_name
+    del inp_model
+    inp_model = None
+    cached_inp_model_name = ''
+    torch.cuda.empty_cache()
+
+
+def set_model_from_name(inp_model_name):
+    global cached_inp_model_name
+    global inp_model
+
+    if inp_model_name == cached_inp_model_name:
+        print (f"Activating Cached Inpaintng Model: {inp_model_name}")
+        return
+
+    print (f"Activating Inpaintng Model: {inp_model_name}")
+    inp_model = inpainting_models[inp_model_name]
+    cached_inp_model_name = inp_model_name
+
+
+def rasg_run(use_painta, prompt, input, seed, eta, negative_prompt, positive_prompt, ddim_steps,
+guidance_scale=7.5, batch_size=4):
+    torch.cuda.empty_cache()
+
+    seed = int(seed)
+    batch_size = max(1, min(int(batch_size), 4))
+
+    image = IImage(input['image']).resize(512)
+    mask = IImage(input['mask']).rgb().resize(512)
+
+    method = ['rasg']
+    if use_painta: method.append('painta')
+
+    inpainted_images = []
+    blended_images = []
+    for i in range(batch_size):
+        inpainted_image = rasg.run(
+            ddim = inp_model,
+            method = '-'.join(method),
+            prompt = prompt,
+            image = image.padx(64),
+            mask = mask.alpha().padx(64),
+            seed = seed+i*1000,
+            eta = eta,
+            prefix = '{}',
+            negative_prompt = negative_prompt,
+            positive_prompt = f', {positive_prompt}',
+            dt = 1000 // ddim_steps,
+            guidance_scale = guidance_scale
+        ).crop(image.size)
+        blended_image = poisson_blend(orig_img = image.data[0], fake_img = inpainted_image.data[0],
+            mask = mask.data[0], dilation = 12)
+        
+        blended_images.append(blended_image)
+        inpainted_images.append(inpainted_image.numpy()[0])
+
+    return blended_images, inpainted_images
+
+
+def sd_run(use_painta, prompt, input, seed, eta, negative_prompt, positive_prompt, ddim_steps,
+guidance_scale=7.5, batch_size=4):
+    torch.cuda.empty_cache()
+
+    seed = int(seed)
+    batch_size = max(1, min(int(batch_size), 4))
+
+    image = IImage(input['image']).resize(512)
+    mask = IImage(input['mask']).rgb().resize(512)
+
+    method = ['default']
+    if use_painta: method.append('painta')
+
+    inpainted_images = []
+    blended_images = []
+    for i in range(batch_size):
+        inpainted_image = sd.run(
+            ddim = inp_model,
+            method = '-'.join(method),
+            prompt = prompt,
+            image = image.padx(64),
+            mask = mask.alpha().padx(64),
+            seed = seed+i*1000,
+            eta = eta,
+            prefix = '{}',
+            negative_prompt = negative_prompt,
+            positive_prompt = f', {positive_prompt}',
+            dt = 1000 // ddim_steps,
+            guidance_scale = guidance_scale
+        ).crop(image.size)
+
+        blended_image = poisson_blend(orig_img = image.data[0], fake_img = inpainted_image.data[0],
+            mask = mask.data[0], dilation = 12)
+
+        blended_images.append(blended_image)
+        inpainted_images.append(inpainted_image.numpy()[0])
+
+    return blended_images, inpainted_images
+
+
+def upscale_run(
+    prompt, input, ddim_steps, seed, use_sam_mask, gallery, img_index,
+negative_prompt='', positive_prompt=', high resolution professional photo'):
+    torch.cuda.empty_cache()
+
+    # Load SR model and SAM predictor
+    # sr_model = models.sd2_sr.load_model()
+    # sam_predictor = None
+    # if use_sam_mask:
+    #     sam_predictor = models.sam.load_model()
+
+    seed = int(seed)
+    img_index = int(img_index)
+ 
+    img_index = 0 if img_index < 0 else img_index
+    img_index = len(gallery) - 1 if img_index >= len(gallery) else img_index
+    img_info = gallery[img_index if img_index >= 0 else 0]
+    inpainted_image = image_from_url_text(img_info)
+    lr_image = IImage(inpainted_image)
+    hr_image = IImage(input['image']).resize(2048)
+    hr_mask = IImage(input['mask']).resize(2048)
+    output_image = sr.run(sr_model, sam_predictor, lr_image, hr_image, hr_mask, prompt=prompt + positive_prompt,
+        noise_level=0, blend_trick=True, blend_output=True, negative_prompt=negative_prompt, 
+        seed=seed, use_sam_mask=use_sam_mask)
+    return output_image.numpy()[0], output_image.numpy()[0]
+
+
+def switch_run(use_rasg, model_name, *args):
+    set_model_from_name(model_name)
+    if use_rasg:
+        return rasg_run(*args)
+    return sd_run(*args)
+
+
+with gr.Blocks(css='style.css') as demo:
+    gr.HTML(
+        """
+        <div style="text-align: center; max-width: 1200px; margin: 20px auto;">
+        <h1 style="font-weight: 900; font-size: 3rem; margin-bottom: 0.5rem">
+            🧑‍🎨 HD-Painter Demo
+        </h1>
+        <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem">
+        Hayk Manukyan<sup>1*</sup>, Andranik Sargsyan<sup>1*</sup>, Barsegh Atanyan<sup>1</sup>, Zhangyang Wang<sup>1,2</sup>, Shant Navasardyan<sup>1</sup>
+        and <a href="https://www.humphreyshi.com/home">Humphrey Shi</a><sup>1,3</sup>
+        </h2>
+        <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem">
+        <sup>1</sup>Picsart AI Resarch (PAIR), <sup>2</sup>UT Austin, <sup>3</sup>Georgia Tech
+        </h2>
+        <h2 style="font-weight: 450; font-size: 1rem; margin: 0rem">
+        [<a href="https://arxiv.org/abs/2312.14091" style="color:blue;">arXiv</a>] 
+        [<a href="https://github.com/Picsart-AI-Research/HD-Painter" style="color:blue;">GitHub</a>]
+        </h2>
+        <h2 style="font-weight: 450; font-size: 1rem; margin: 0.7rem auto; max-width: 1000px">
+        <b>HD-Painter</b> enables prompt-faithfull and high resolution (up to 2k) image inpainting upon any diffusion-based image inpainting method.
+        </h2>
+        </div>
+        """)
+
+    if on_huggingspace:
+        gr.HTML("""
+        <p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
+        <br/>
+        <a href="https://huggingface.co/spaces/PAIR/HD-Painter?duplicate=true">
+        <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
+        </p>""")
+
+    with open('script.js', 'r') as f:
+        js_str = f.read()
+
+    demo.load(_js=js_str)
+
+    with gr.Row():
+        with gr.Column():
+            model_picker = gr.Dropdown(
+                list(inpainting_models.keys()),
+                value=0,
+                label = "Please select a model!",
+            )
+        with gr.Column():
+            use_painta = gr.Checkbox(value = True, label = "Use PAIntA")
+            use_rasg = gr.Checkbox(value = True, label = "Use RASG")
+
+    prompt = gr.Textbox(label = "Inpainting Prompt")
+    with gr.Row():
+        with gr.Column():
+            input = gr.ImageMask(label = "Input Image", brush_color='#ff0000', elem_id="inputmask")
+            
+            with gr.Row():
+                inpaint_btn = gr.Button("Inpaint", scale = 0)
+   
+            with gr.Accordion('Advanced options', open=False):
+                guidance_scale = gr.Slider(minimum = 0, maximum = 30, value = 7.5, label = "Guidance Scale")
+                eta = gr.Slider(minimum = 0, maximum = 1, value = 0.1, label = "eta")
+                ddim_steps = gr.Slider(minimum = 10, maximum = 100, value = 50, step =  1, label = 'Number of diffusion steps')
+                with gr.Row():
+                    seed = gr.Number(value = 49123, label = "Seed")
+                    batch_size = gr.Number(value = 1, label = "Batch size", minimum=1, maximum=4) 
+                negative_prompt = gr.Textbox(value=negative_prompt_str, label = "Negative prompt", lines=3)
+                positive_prompt = gr.Textbox(value=positive_prompt_str, label = "Positive prompt", lines=1)
+
+        with gr.Column():
+            with gr.Row():
+                output_gallery = gr.Gallery(
+                    [],
+                    columns = 4,
+                    preview = True,
+                    allow_preview = True,
+                    object_fit='scale-down',
+                    elem_id='outputgallery'
+                )
+            with gr.Row():
+                upscale_btn = gr.Button("Send to Inpainting-Specialized Super-Resolution (x4)", scale = 1)
+            with gr.Row():
+                use_sam_mask = gr.Checkbox(value = False, label = "Use SAM mask for background preservation (for SR only, experimental feature)")
+            with gr.Row():
+                hires_image = gr.Image(label = "Hi-res Image")
+    
+    label = gr.Markdown("## High-Resolution Generation Samples (2048px large side)")
+    
+    with gr.Column():
+        example_container = gr.Gallery(
+            example_previews,
+            columns = 4,
+            preview = True,
+            allow_preview = True,
+            object_fit='scale-down'
+        )
+
+        gr.Examples(
+            [
+                example_inputs[i] + [[example_previews[i]]]
+                for i in range(len(example_previews))
+            ],
+            [input, prompt, example_container]
+        )
+
+    mock_output_gallery = gr.Gallery([], columns = 4, visible=False)
+    mock_hires = gr.Image(label = "__MHRO__", visible = False)
+    html_info = gr.HTML(elem_id=f'html_info', elem_classes="infotext")
+
+    inpaint_btn.click(
+        fn=switch_run, 
+        inputs=[
+            use_rasg,
+            model_picker,
+            use_painta,
+            prompt,
+            input,
+            seed,
+            eta,
+            negative_prompt,
+            positive_prompt,
+            ddim_steps,
+            guidance_scale,
+            batch_size
+        ], 
+        outputs=[output_gallery, mock_output_gallery], 
+        api_name="inpaint"
+    )
+    upscale_btn.click(
+        fn=upscale_run, 
+        inputs=[
+            prompt, 
+            input,
+            ddim_steps,
+            seed,
+            use_sam_mask,
+            mock_output_gallery,
+            html_info
+        ],
+        outputs=[hires_image, mock_hires], 
+        api_name="upscale",
+        _js="function(a, b, c, d, e, f, g){ return [a, b, c, d, e, f, selected_gallery_index()] }",
+    )
+
+demo.queue()
+demo.launch(share=True, allowed_paths=[TMP_DIR])

assets/.gitignore

@@ -0,0 +1 @@
+models/

assets/config/ddpm/v1.yaml

@@ -0,0 +1,14 @@
+linear_start: 0.00085
+linear_end: 0.0120
+num_timesteps_cond: 1
+log_every_t: 200
+timesteps: 1000
+first_stage_key: "jpg"
+cond_stage_key: "txt"
+image_size: 64
+channels: 4
+cond_stage_trainable: false
+conditioning_key: crossattn
+monitor: val/loss_simple_ema
+scale_factor: 0.18215
+use_ema: False # we set this to false because this is an inference only config

assets/config/ddpm/v2-upsample.yaml

@@ -0,0 +1,24 @@
+parameterization: "v"
+low_scale_key: "lr"
+linear_start: 0.0001
+linear_end: 0.02
+num_timesteps_cond: 1
+log_every_t: 200
+timesteps: 1000
+first_stage_key: "jpg"
+cond_stage_key: "txt"
+image_size: 128
+channels: 4
+cond_stage_trainable: false
+conditioning_key: "hybrid-adm"
+monitor: val/loss_simple_ema
+scale_factor: 0.08333
+use_ema: False
+
+low_scale_config:
+  target: ldm.modules.diffusionmodules.upscaling.ImageConcatWithNoiseAugmentation
+  params:
+    noise_schedule_config: # image space
+      linear_start: 0.0001
+      linear_end: 0.02
+    max_noise_level: 350

assets/config/encoders/clip.yaml

@@ -0,0 +1 @@
+__class__: smplfusion.models.encoders.clip_embedder.FrozenCLIPEmbedder

assets/config/encoders/openclip.yaml

@@ -0,0 +1,4 @@
+__class__: smplfusion.models.encoders.open_clip_embedder.FrozenOpenCLIPEmbedder
+__init__:
+  freeze: True
+  layer: "penultimate"

assets/config/unet/inpainting/v1.yaml

@@ -0,0 +1,15 @@
+__class__: smplfusion.models.unet.UNetModel
+__init__:
+  image_size: 32 # unused
+  in_channels: 9  # 4 data + 4 downscaled image + 1 mask
+  out_channels: 4
+  model_channels: 320
+  attention_resolutions: [ 4, 2, 1 ]
+  num_res_blocks: 2
+  channel_mult: [ 1, 2, 4, 4 ]
+  num_heads: 8
+  use_spatial_transformer: True
+  transformer_depth: 1
+  context_dim: 768
+  use_checkpoint: False
+  legacy: False

assets/config/unet/inpainting/v2.yaml

@@ -0,0 +1,16 @@
+__class__: smplfusion.models.unet.UNetModel
+__init__:
+  use_checkpoint: False
+  image_size: 32 # unused
+  in_channels: 9
+  out_channels: 4
+  model_channels: 320
+  attention_resolutions: [ 4, 2, 1 ]
+  num_res_blocks: 2
+  channel_mult: [ 1, 2, 4, 4 ]
+  num_head_channels: 64 # need to fix for flash-attn
+  use_spatial_transformer: True
+  use_linear_in_transformer: True
+  transformer_depth: 1
+  context_dim: 1024
+  legacy: False

assets/config/unet/upsample/v2.yaml

@@ -0,0 +1,19 @@
+__class__: smplfusion.models.unet.UNetModel
+__init__:
+  use_checkpoint: False
+  num_classes: 1000  # timesteps for noise conditioning (here constant, just need one)
+  image_size: 128
+  in_channels: 7
+  out_channels: 4
+  model_channels: 256
+  attention_resolutions: [ 2,4,8]
+  num_res_blocks: 2
+  channel_mult: [ 1, 2, 2, 4]
+  disable_self_attentions: [True, True, True, False]
+  disable_middle_self_attn: False
+  num_heads: 8
+  use_spatial_transformer: True
+  transformer_depth: 1
+  context_dim: 1024
+  legacy: False
+  use_linear_in_transformer: True

assets/config/vae-upsample.yaml

@@ -0,0 +1,16 @@
+__class__: smplfusion.models.vae.AutoencoderKL
+__init__:
+  embed_dim: 4
+  ddconfig:
+    double_z: True
+    z_channels: 4
+    resolution: 256
+    in_channels: 3
+    out_ch: 3
+    ch: 128
+    ch_mult: [ 1,2,4 ]
+    num_res_blocks: 2
+    attn_resolutions: [ ]
+    dropout: 0.0
+  lossconfig:
+    target: torch.nn.Identity

assets/config/vae.yaml

@@ -0,0 +1,17 @@
+__class__: smplfusion.models.vae.AutoencoderKL
+__init__:
+  embed_dim: 4
+  monitor: val/rec_loss
+  ddconfig:
+    double_z: true
+    z_channels: 4
+    resolution: 256
+    in_channels: 3
+    out_ch: 3
+    ch: 128
+    ch_mult: [1,2,4,4]
+    num_res_blocks: 2
+    attn_resolutions: []
+    dropout: 0.0
+  lossconfig:
+    target: torch.nn.Identity

lib/methods/rasg.py

@@ -0,0 +1,88 @@
+import torch
+from lib.utils.iimage import IImage
+from pytorch_lightning import seed_everything
+from tqdm import tqdm
+
+from lib.smplfusion import share, router, attentionpatch, transformerpatch
+from lib.smplfusion.patches.attentionpatch import painta
+from lib.utils import tokenize, scores
+
+verbose = False
+
+
+def init_painta(token_idx):
+    # Initialize painta
+    router.attention_forward = attentionpatch.painta.forward
+    router.basic_transformer_forward = transformerpatch.painta.forward
+    painta.painta_on = True
+    painta.painta_res = [16, 32]
+    painta.token_idx = token_idx
+
+def init_guidance():
+    # Setup model for guidance only!
+    router.attention_forward = attentionpatch.default.forward_and_save
+    router.basic_transformer_forward = transformerpatch.default.forward
+
+def run(ddim, method, prompt, image, mask, seed, eta, prefix, negative_prompt, positive_prompt, dt, guidance_scale):
+    # Text condition
+    prompt = prefix.format(prompt)
+    context = ddim.encoder.encode([negative_prompt, prompt + positive_prompt])
+    token_idx = list(range(1 + prefix.split(' ').index('{}'), tokenize(prompt).index('<end_of_text>')))
+    token_idx += [tokenize(prompt + positive_prompt).index('<end_of_text>')]
+    
+    # Initialize painta
+    if 'painta' in method: init_painta(token_idx)
+    else: init_guidance()
+
+    # Image condition
+    unet_condition = ddim.get_inpainting_condition(image, mask)
+    share.set_mask(mask)
+
+    # Starting latent
+    seed_everything(seed)
+    zt = torch.randn((1,4) + unet_condition.shape[2:]).cuda()
+
+    # Setup unet for guidance
+    ddim.unet.requires_grad_(True)
+    
+    pbar = tqdm(range(999, 0, -dt)) if verbose else range(999, 0, -dt)
+
+    for timestep in share.DDIMIterator(pbar):
+        if 'painta' in method and share.timestep <= 500: init_guidance()
+        
+        zt = zt.detach()
+        zt.requires_grad = True
+
+        # Reset storage
+        share._crossattn_similarity_res16 = []
+
+        # Run the model
+        _zt = zt if unet_condition is None else torch.cat([zt, unet_condition], 1)
+        eps_uncond, eps = ddim.unet(
+            torch.cat([_zt, _zt]), 
+            timesteps = torch.tensor([timestep, timestep]).cuda(), 
+            context = context
+        ).detach().chunk(2)
+        
+        # Unconditional guidance
+        eps = (eps_uncond + guidance_scale * (eps - eps_uncond))
+        z0 = (zt - share.schedule.sqrt_one_minus_alphas[timestep] * eps) / share.schedule.sqrt_alphas[timestep]
+
+        # Gradient Computation
+        score = scores.bce(share._crossattn_similarity_res16, share.mask16, token_idx = token_idx)
+        score.backward()
+        grad = zt.grad.detach()
+        ddim.unet.zero_grad() # Cleanup already
+
+        # DDIM Step
+        with torch.no_grad():
+            sigma = share.schedule.sigma(share.timestep, dt)
+            # Standartization
+            grad -= grad.mean()
+            grad /= grad.std()
+            
+            zt = share.schedule.sqrt_alphas[share.timestep - dt] * z0 + torch.sqrt(1 - share.schedule.alphas[share.timestep - dt] - sigma ** 2) * eps + eta * sigma * grad
+
+    with torch.no_grad():
+        output_image = IImage(ddim.vae.decode(z0 / ddim.config.scale_factor))
+    return output_image

lib/methods/sd.py

@@ -0,0 +1,74 @@
+import torch
+from pytorch_lightning import seed_everything
+from tqdm import tqdm
+
+from lib.utils.iimage import IImage
+from lib.smplfusion import share, router, attentionpatch, transformerpatch
+from lib.smplfusion.patches.attentionpatch import painta
+from lib.utils import tokenize
+
+verbose = False
+
+
+def init_painta(token_idx):
+    # Initialize painta
+    router.attention_forward = attentionpatch.painta.forward
+    router.basic_transformer_forward = transformerpatch.painta.forward
+    painta.painta_on = True
+    painta.painta_res = [16, 32]
+    painta.token_idx = token_idx
+
+def run(
+    ddim,
+    method,
+    prompt,
+    image,
+    mask,
+    seed,
+    eta,
+    prefix,
+    negative_prompt,
+    positive_prompt,
+    dt,
+    guidance_scale
+):
+    # Text condition
+    context = ddim.encoder.encode([negative_prompt, prompt + positive_prompt])
+    token_idx = list(range(1 + prefix.split(' ').index('{}'), tokenize(prompt).index('<end_of_text>')))
+    token_idx += [tokenize(prompt + positive_prompt).index('<end_of_text>')]
+
+    # Setup painta if needed
+    if 'painta' in method: init_painta(token_idx)
+    else: router.reset() 
+    
+    # Image condition    
+    unet_condition = ddim.get_inpainting_condition(image, mask)
+    share.set_mask(mask)
+    
+    # Starting latent
+    seed_everything(seed)
+    zt = torch.randn((1,4) + unet_condition.shape[2:]).cuda()
+
+    # Turn off gradients
+    ddim.unet.requires_grad_(False)
+    
+    pbar = tqdm(range(999, 0, -dt)) if verbose else range(999, 0, -dt)
+
+    for timestep in share.DDIMIterator(pbar):
+        if share.timestep <= 500: router.reset()
+
+        _zt = zt if unet_condition is None else torch.cat([zt, unet_condition], 1)
+        eps_uncond, eps = ddim.unet(
+            torch.cat([_zt, _zt]), 
+            timesteps = torch.tensor([timestep, timestep]).cuda(), 
+            context = context
+        ).chunk(2)
+
+        eps = (eps_uncond + guidance_scale * (eps - eps_uncond))
+        z0 = (zt - share.schedule.sqrt_one_minus_alphas[timestep] * eps) / share.schedule.sqrt_alphas[timestep]
+        zt = share.schedule.sqrt_alphas[timestep - dt] * z0 + share.schedule.sqrt_one_minus_alphas[timestep - dt] * eps
+
+    with torch.no_grad():
+        output_image = IImage(ddim.vae.decode(z0 / ddim.config.scale_factor))
+    
+    return output_image

lib/methods/sr.py

@@ -0,0 +1,141 @@
+import os
+from functools import partial
+from glob import glob
+from pathlib import Path as PythonPath
+
+import cv2
+import torchvision.transforms.functional as TvF
+import torch
+import torch.nn as nn
+import numpy as np
+from inspect import isfunction
+from PIL import Image
+
+from lib import smplfusion
+from lib.smplfusion import share, router, attentionpatch, transformerpatch
+from lib.utils.iimage import IImage
+from lib.utils import poisson_blend
+from lib.models.sd2_sr import predict_eps_from_z_and_v, predict_start_from_z_and_v
+
+
+def refine_mask(hr_image, hr_mask, lr_image, sam_predictor):
+    lr_mask = hr_mask.resize(512)
+
+    x_min, y_min, rect_w, rect_h = cv2.boundingRect(lr_mask.data[0][:, :, 0])
+    x_min = max(x_min - 1, 0)
+    y_min = max(y_min - 1, 0)
+    x_max = x_min + rect_w + 1
+    y_max = y_min + rect_h + 1
+
+    input_box = np.array([x_min, y_min, x_max, y_max])
+
+    sam_predictor.set_image(hr_image.resize(512).data[0])
+    masks, _, _ = sam_predictor.predict(
+        point_coords=None,
+        point_labels=None,
+        box=input_box[None, :],
+        multimask_output=True,
+    )
+    dilation_kernel = np.ones((13, 13))
+    original_object_mask = (np.sum(masks, axis=0) > 0).astype(np.uint8)
+    original_object_mask = cv2.dilate(original_object_mask, dilation_kernel)
+
+    sam_predictor.set_image(lr_image.resize(512).data[0])
+    masks, _, _ = sam_predictor.predict(
+        point_coords=None,
+        point_labels=None,
+        box=input_box[None, :],
+        multimask_output=True,
+    )
+    dilation_kernel = np.ones((3, 3))
+    inpainted_object_mask = (np.sum(masks, axis=0) > 0).astype(np.uint8)
+    inpainted_object_mask = cv2.dilate(inpainted_object_mask, dilation_kernel)
+
+    lr_mask_masking = ((original_object_mask + inpainted_object_mask ) > 0).astype(np.uint8)
+    new_mask = lr_mask.data[0] * lr_mask_masking[:, :, np.newaxis]
+    new_mask = IImage(new_mask).resize(2048, resample = Image.BICUBIC)
+    return new_mask
+
+
+def run(ddim, sam_predictor, lr_image, hr_image, hr_mask, prompt = 'high resolution professional photo', noise_level=20,
+blend_output = True, blend_trick = True, no_superres = False,
+dt = 20, seed = 1, guidance_scale = 7.5, negative_prompt = '', use_sam_mask = False, dtype=torch.bfloat16):
+    torch.manual_seed(seed)
+
+    router.attention_forward = attentionpatch.default.forward_xformers
+    router.basic_transformer_forward = transformerpatch.default.forward
+
+    if use_sam_mask:
+        with torch.no_grad():
+            hr_mask = refine_mask(hr_image, hr_mask, lr_image, sam_predictor)
+
+    orig_h, orig_w = hr_image.torch().shape[2], hr_image.torch().shape[3]
+    hr_image = hr_image.padx(256, padding_mode='reflect')
+    hr_mask = hr_mask.padx(256, padding_mode='reflect').dilate(19)
+    hr_mask_orig = hr_mask
+    lr_image = lr_image.padx(64, padding_mode='reflect')
+    lr_mask = hr_mask.resize((lr_image.torch().shape[2], lr_image.torch().shape[3]), resample = Image.BICUBIC).alpha().torch(vmin=0).cuda()
+    lr_mask = TvF.gaussian_blur(lr_mask, kernel_size=19)
+
+    if no_superres:
+        output_tensor = lr_image.resize((hr_image.torch().shape[2], hr_image.torch().shape[3]), resample = Image.BICUBIC).torch().cuda()
+        output_tensor = (255*((output_tensor.clip(-1, 1) + 1) / 2)).to(torch.uint8)
+        output_tensor = poisson_blend(
+            orig_img=hr_image.data[0][:orig_h, :orig_w, :],
+            fake_img=output_tensor.cpu().permute(0, 2, 3, 1)[0].numpy()[:orig_h, :orig_w, :],
+            mask=hr_mask_orig.alpha().data[0][:orig_h, :orig_w, :]
+        )
+        return IImage(output_tensor[:orig_h, :orig_w, :])
+                                     
+    # encode hr image
+    with torch.no_grad():
+        hr_z0 = ddim.vae.encode(hr_image.torch().cuda().to(dtype)).mean * ddim.config.scale_factor
+
+    assert hr_z0.shape[2] == lr_image.torch().shape[2]
+    assert hr_z0.shape[3] == lr_image.torch().shape[3]                                  
+                                     
+    unet_condition = lr_image.cuda().torch().to(memory_format=torch.contiguous_format).to(dtype)
+    zT = torch.randn((1,4,unet_condition.shape[2], unet_condition.shape[3])).cuda().to(dtype)
+    
+    with torch.no_grad():
+        context = ddim.encoder.encode([negative_prompt, prompt])
+
+    noise_level = torch.Tensor(1 * [noise_level]).to('cuda').long()
+    unet_condition, noise_level = ddim.low_scale_model(unet_condition, noise_level=noise_level)
+
+    with torch.autocast('cuda'), torch.no_grad():
+        zt = zT
+        for index,t in enumerate(range(999, 0, -dt)):
+           
+            _zt = zt if unet_condition is None else torch.cat([zt, unet_condition], 1)
+
+            eps_uncond, eps = ddim.unet(
+                torch.cat([_zt, _zt]).to(dtype), 
+                timesteps = torch.tensor([t, t]).cuda(), 
+                context = context,
+                y=torch.cat([noise_level]*2)
+            ).chunk(2)
+
+            ts = torch.full((zt.shape[0],), t, device='cuda', dtype=torch.long)
+            model_output = (eps_uncond + guidance_scale * (eps - eps_uncond))
+            eps = predict_eps_from_z_and_v(ddim.schedule, zt, ts, model_output).to(dtype)
+            z0 = predict_start_from_z_and_v(ddim.schedule, zt, ts, model_output).to(dtype)
+
+            if blend_trick:
+                z0 = z0 * lr_mask + hr_z0 * (1-lr_mask)
+            
+            zt = ddim.schedule.sqrt_alphas[t - dt] * z0 + ddim.schedule.sqrt_one_minus_alphas[t - dt] * eps
+
+    with torch.no_grad():
+        output_tensor = ddim.vae.decode(z0.to(dtype) / ddim.config.scale_factor)
+
+    if blend_output:
+        output_tensor = (255*((output_tensor + 1) / 2).clip(0, 1)).to(torch.uint8)
+        output_tensor = poisson_blend(
+            orig_img=hr_image.data[0][:orig_h, :orig_w, :],
+            fake_img=output_tensor.cpu().permute(0, 2, 3, 1)[0].numpy()[:orig_h, :orig_w, :],
+            mask=hr_mask_orig.alpha().data[0][:orig_h, :orig_w, :]
+        )
+        return IImage(output_tensor[:orig_h, :orig_w, :])
+    else:
+        return IImage(output_tensor[:, :, :orig_h, :orig_w])      

lib/models/__init__.py

@@ -0,0 +1 @@
+from . import sd2_inp, ds_inp, sd15_inp, sd2_sr, sam

lib/models/common.py

@@ -0,0 +1,49 @@
+import importlib
+import requests
+from pathlib import Path
+from os.path import dirname
+
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+
+PROJECT_DIR = dirname(dirname(dirname(__file__)))
+CONFIG_FOLDER =  f'{PROJECT_DIR}/assets/config'
+MODEL_FOLDER =  f'{PROJECT_DIR}/assets/models'
+
+
+def download_file(url, save_path, chunk_size=1024):
+    try:
+        save_path = Path(save_path)
+        if save_path.exists():
+            print(f'{save_path.name} exists')
+            return
+        save_path.parent.mkdir(exist_ok=True, parents=True)
+        resp = requests.get(url, stream=True)
+        total = int(resp.headers.get('content-length', 0))
+        with open(save_path, 'wb') as file, tqdm(
+            desc=save_path.name,
+            total=total,
+            unit='iB',
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar:
+            for data in resp.iter_content(chunk_size=chunk_size):
+                size = file.write(data)
+                bar.update(size)
+        print(f'{save_path.name} download finished')
+    except Exception as e:
+        raise Exception(f"Download failed: {e}")
+
+
+def get_obj_from_str(string):
+    module, cls = string.rsplit(".", 1)
+    try:
+        return getattr(importlib.import_module(module, package=None), cls)
+    except:
+        return getattr(importlib.import_module('lib.' + module, package=None), cls)
+
+
+def load_obj(path):
+    objyaml = OmegaConf.load(path)
+    return get_obj_from_str(objyaml['__class__'])(**objyaml.get("__init__", {}))

lib/models/ds_inp.py

@@ -0,0 +1,46 @@
+import importlib
+from omegaconf import OmegaConf
+import torch
+import safetensors
+import safetensors.torch
+
+from lib.smplfusion import DDIM, share, scheduler
+from .common import *
+
+
+MODEL_PATH = f'{MODEL_FOLDER}/dreamshaper/dreamshaper_8Inpainting.safetensors'
+DOWNLOAD_URL = 'https://civitai.com/api/download/models/131004'
+
+# pre-download
+download_file(DOWNLOAD_URL, MODEL_PATH)
+
+
+def load_model():
+    print ("Loading model: Dreamshaper Inpainting V8")
+
+    download_file(DOWNLOAD_URL, MODEL_PATH)
+    
+    state_dict = safetensors.torch.load_file(MODEL_PATH)
+
+    config = OmegaConf.load(f'{CONFIG_FOLDER}/ddpm/v1.yaml')
+    unet = load_obj(f'{CONFIG_FOLDER}/unet/inpainting/v1.yaml').eval().cuda()
+    vae = load_obj(f'{CONFIG_FOLDER}/vae.yaml').eval().cuda()
+    encoder = load_obj(f'{CONFIG_FOLDER}/encoders/clip.yaml').eval().cuda()
+
+    extract = lambda state_dict, model: {x[len(model)+1:]:y for x,y in state_dict.items() if model in x}
+    unet_state = extract(state_dict, 'model.diffusion_model')
+    encoder_state = extract(state_dict, 'cond_stage_model')
+    vae_state = extract(state_dict, 'first_stage_model')
+
+    unet.load_state_dict(unet_state)
+    encoder.load_state_dict(encoder_state)
+    vae.load_state_dict(vae_state)
+
+    unet = unet.requires_grad_(False)
+    encoder = encoder.requires_grad_(False)
+    vae = vae.requires_grad_(False)
+
+    ddim = DDIM(config, vae, encoder, unet)
+    share.schedule = scheduler.linear(config.timesteps, config.linear_start, config.linear_end)
+
+    return ddim

lib/models/sam.py

@@ -0,0 +1,20 @@
+from segment_anything import sam_model_registry, SamPredictor
+from .common import *
+
+MODEL_PATH = f'{MODEL_FOLDER}/sam/sam_vit_h_4b8939.pth'
+DOWNLOAD_URL = 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth'
+
+# pre-download
+download_file(DOWNLOAD_URL, MODEL_PATH)
+
+
+def load_model():
+    print ("Loading model: SAM")
+    download_file(DOWNLOAD_URL, MODEL_PATH)
+    model_type = "vit_h"
+    device = "cuda"
+    sam = sam_model_registry[model_type](checkpoint=MODEL_PATH)
+    sam.to(device=device)
+    sam_predictor = SamPredictor(sam)
+    print ("SAM loaded")
+    return sam_predictor

lib/models/sd15_inp.py

@@ -0,0 +1,44 @@
+from omegaconf import OmegaConf
+import torch
+
+from lib.smplfusion import DDIM, share, scheduler
+from .common import *
+
+
+DOWNLOAD_URL = 'https://huggingface.co/runwayml/stable-diffusion-inpainting/resolve/main/sd-v1-5-inpainting.ckpt?download=true'
+MODEL_PATH = f'{MODEL_FOLDER}/sd-1-5-inpainting/sd-v1-5-inpainting.ckpt'
+
+# pre-download
+download_file(DOWNLOAD_URL, MODEL_PATH)
+
+
+def load_model():
+    download_file(DOWNLOAD_URL, MODEL_PATH)
+
+    state_dict = torch.load(MODEL_PATH)['state_dict']
+
+    config = OmegaConf.load(f'{CONFIG_FOLDER}/ddpm/v1.yaml')
+
+    print ("Loading model: Stable-Inpainting 1.5")
+
+    unet = load_obj(f'{CONFIG_FOLDER}/unet/inpainting/v1.yaml').eval().cuda()
+    vae = load_obj(f'{CONFIG_FOLDER}/vae.yaml').eval().cuda()
+    encoder = load_obj(f'{CONFIG_FOLDER}/encoders/clip.yaml').eval().cuda()
+
+    extract = lambda state_dict, model: {x[len(model)+1:]:y for x,y in state_dict.items() if model in x}
+    unet_state = extract(state_dict, 'model.diffusion_model')
+    encoder_state = extract(state_dict, 'cond_stage_model')
+    vae_state = extract(state_dict, 'first_stage_model')
+
+    unet.load_state_dict(unet_state)
+    encoder.load_state_dict(encoder_state)
+    vae.load_state_dict(vae_state)
+
+    unet = unet.requires_grad_(False)
+    encoder = encoder.requires_grad_(False)
+    vae = vae.requires_grad_(False)
+
+    ddim = DDIM(config, vae, encoder, unet)
+    share.schedule = scheduler.linear(config.timesteps, config.linear_start, config.linear_end)
+
+    return ddim

lib/models/sd2_inp.py

@@ -0,0 +1,47 @@
+import safetensors
+import safetensors.torch
+import torch
+from omegaconf import OmegaConf
+
+from lib.smplfusion import DDIM, share, scheduler
+from .common import *
+
+MODEL_PATH = f'{MODEL_FOLDER}/sd-2-0-inpainting/512-inpainting-ema.safetensors'
+DOWNLOAD_URL = 'https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/resolve/main/512-inpainting-ema.safetensors?download=true'
+
+# pre-download
+download_file(DOWNLOAD_URL, MODEL_PATH)
+
+
+def load_model():
+    print ("Loading model: Stable-Inpainting 2.0")
+
+    download_file(DOWNLOAD_URL, MODEL_PATH)
+
+    state_dict = safetensors.torch.load_file(MODEL_PATH)
+
+    config = OmegaConf.load(f'{CONFIG_FOLDER}/ddpm/v1.yaml')
+
+    unet = load_obj(f'{CONFIG_FOLDER}/unet/inpainting/v2.yaml').eval().cuda()
+    vae = load_obj(f'{CONFIG_FOLDER}/vae.yaml').eval().cuda()
+    encoder = load_obj(f'{CONFIG_FOLDER}/encoders/openclip.yaml').eval().cuda()
+    ddim = DDIM(config, vae, encoder, unet)
+
+    extract = lambda state_dict, model: {x[len(model)+1:]:y for x,y in state_dict.items() if model in x}
+    unet_state = extract(state_dict, 'model.diffusion_model')
+    encoder_state = extract(state_dict, 'cond_stage_model')
+    vae_state = extract(state_dict, 'first_stage_model')
+
+    unet.load_state_dict(unet_state)
+    encoder.load_state_dict(encoder_state)
+    vae.load_state_dict(vae_state)
+
+    unet = unet.requires_grad_(False)
+    encoder = encoder.requires_grad_(False)
+    vae = vae.requires_grad_(False)
+
+    ddim = DDIM(config, vae, encoder, unet)
+    share.schedule = scheduler.linear(config.timesteps, config.linear_start, config.linear_end)
+
+    print('Stable-Inpainting 2.0 loaded')
+    return ddim

lib/models/sd2_sr.py

@@ -0,0 +1,204 @@
+import importlib
+from functools import partial
+
+import cv2
+import numpy as np
+import safetensors
+import safetensors.torch
+import torch
+import torch.nn as nn
+from inspect import isfunction
+from omegaconf import OmegaConf
+
+from lib.smplfusion import DDIM, share, scheduler
+from .common import *
+
+
+DOWNLOAD_URL = 'https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler/resolve/main/x4-upscaler-ema.safetensors?download=true'
+MODEL_PATH = f'{MODEL_FOLDER}/sd-2-0-upsample/x4-upscaler-ema.safetensors'
+
+# pre-download
+download_file(DOWNLOAD_URL, MODEL_PATH)
+
+
+def exists(x):
+    return x is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def predict_eps_from_z_and_v(schedule, x_t, t, v):
+    return (
+        extract_into_tensor(schedule.sqrt_alphas.cuda(), t, x_t.shape) * v +
+        extract_into_tensor(schedule.sqrt_one_minus_alphas.cuda(), t, x_t.shape) * x_t
+    )
+
+
+def predict_start_from_z_and_v(schedule, x_t, t, v):
+    return (
+        extract_into_tensor(schedule.sqrt_alphas.cuda(), t, x_t.shape) * x_t -
+        extract_into_tensor(schedule.sqrt_one_minus_alphas.cuda(), t, x_t.shape) * v
+    )
+
+
+def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+    if schedule == "linear":
+        betas = (
+            torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
+        )
+
+    elif schedule == "cosine":
+        timesteps = (
+            torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s
+        )
+        alphas = timesteps / (1 + cosine_s) * np.pi / 2
+        alphas = torch.cos(alphas).pow(2)
+        alphas = alphas / alphas[0]
+        betas = 1 - alphas[1:] / alphas[:-1]
+        betas = np.clip(betas, a_min=0, a_max=0.999)
+
+    elif schedule == "sqrt_linear":
+        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
+    elif schedule == "sqrt":
+        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
+    else:
+        raise ValueError(f"schedule '{schedule}' unknown.")
+    return betas.numpy()
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+class AbstractLowScaleModel(nn.Module):
+    # for concatenating a downsampled image to the latent representation
+    def __init__(self, noise_schedule_config=None):
+        super(AbstractLowScaleModel, self).__init__()
+        if noise_schedule_config is not None:
+            self.register_schedule(**noise_schedule_config)
+
+    def register_schedule(self, beta_schedule="linear", timesteps=1000,
+                          linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+        betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
+                                   cosine_s=cosine_s)
+        alphas = 1. - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+
+        timesteps, = betas.shape
+        self.num_timesteps = int(timesteps)
+        self.linear_start = linear_start
+        self.linear_end = linear_end
+        assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
+
+        to_torch = partial(torch.tensor, dtype=torch.float32)
+
+        self.register_buffer('betas', to_torch(betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
+
+    def q_sample(self, x_start, t, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+                extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
+
+    def forward(self, x):
+        return x, None
+
+    def decode(self, x):
+        return x
+
+
+class ImageConcatWithNoiseAugmentation(AbstractLowScaleModel):
+    def __init__(self, noise_schedule_config, max_noise_level=1000, to_cuda=False):
+        super().__init__(noise_schedule_config=noise_schedule_config)
+        self.max_noise_level = max_noise_level
+
+    def forward(self, x, noise_level=None):
+        if noise_level is None:
+            noise_level = torch.randint(0, self.max_noise_level, (x.shape[0],), device=x.device).long()
+        else:
+            assert isinstance(noise_level, torch.Tensor)
+        z = self.q_sample(x, noise_level)
+        return z, noise_level
+
+
+def get_obj_from_str(string):
+    module, cls = string.rsplit(".", 1)
+    try:
+        return getattr(importlib.import_module(module, package=None), cls)
+    except:
+        return getattr(importlib.import_module('lib.' + module, package=None), cls)
+def load_obj(path):
+    objyaml = OmegaConf.load(path)
+    return get_obj_from_str(objyaml['__class__'])(**objyaml.get("__init__", {}))
+    
+
+def load_model(dtype=torch.bfloat16):
+    print ("Loading model: SD2 superresolution...")
+
+    download_file(DOWNLOAD_URL, MODEL_PATH)
+
+    state_dict = safetensors.torch.load_file(MODEL_PATH)
+
+    config = OmegaConf.load(f'{CONFIG_FOLDER}/ddpm/v2-upsample.yaml')
+
+    unet = load_obj(f'{CONFIG_FOLDER}/unet/upsample/v2.yaml').eval().cuda()
+    vae = load_obj(f'{CONFIG_FOLDER}/vae-upsample.yaml').eval().cuda()
+    encoder = load_obj(f'{CONFIG_FOLDER}/encoders/openclip.yaml').eval().cuda()
+    ddim = DDIM(config, vae, encoder, unet)
+
+    extract = lambda state_dict, model: {x[len(model)+1:]:y for x,y in state_dict.items() if model in x}
+    unet_state = extract(state_dict, 'model.diffusion_model')
+    encoder_state = extract(state_dict, 'cond_stage_model')
+    vae_state = extract(state_dict, 'first_stage_model')
+
+    unet.load_state_dict(unet_state)
+    encoder.load_state_dict(encoder_state)
+    vae.load_state_dict(vae_state)
+
+    unet = unet.requires_grad_(False)
+    encoder = encoder.requires_grad_(False)
+    vae = vae.requires_grad_(False)
+    
+    unet.to(dtype)
+    vae.to(dtype)
+    encoder.to(dtype)
+
+    ddim = DDIM(config, vae, encoder, unet)
+
+    params = {
+        'noise_schedule_config': {
+            'linear_start': 0.0001,
+            'linear_end': 0.02
+        },
+        'max_noise_level': 350
+    }
+
+    low_scale_model = ImageConcatWithNoiseAugmentation(**params).eval().to('cuda')
+    low_scale_model.train = disabled_train
+    for param in low_scale_model.parameters():
+        param.requires_grad = False
+
+    ddim.low_scale_model = low_scale_model
+    print('SD2 superresolution loaded')
+    return ddim