libcom 0.1.8__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (310) hide show
  1. libcom/__init__.py +23 -0
  2. libcom/color_transfer/__init__.py +3 -0
  3. libcom/color_transfer/reinhard.py +57 -0
  4. libcom/fopa_heat_map/__init__.py +4 -0
  5. libcom/fopa_heat_map/fopa_heat_map.py +212 -0
  6. libcom/fopa_heat_map/source/__init__.py +0 -0
  7. libcom/fopa_heat_map/source/backbone/ResNet.py +229 -0
  8. libcom/fopa_heat_map/source/backbone/__init__.py +0 -0
  9. libcom/fopa_heat_map/source/data/OBdataset.py +292 -0
  10. libcom/fopa_heat_map/source/data/__init__.py +0 -0
  11. libcom/fopa_heat_map/source/data/all_transforms.py +25 -0
  12. libcom/fopa_heat_map/source/network/BaseBlocks.py +26 -0
  13. libcom/fopa_heat_map/source/network/DynamicModules.py +43 -0
  14. libcom/fopa_heat_map/source/network/ObPlaNet_simple.py +112 -0
  15. libcom/fopa_heat_map/source/network/__init__.py +1 -0
  16. libcom/fopa_heat_map/source/network/tensor_ops.py +37 -0
  17. libcom/fopa_heat_map/source/prepare_multi_fg_scales.py +85 -0
  18. libcom/fos_score/__init__.py +4 -0
  19. libcom/fos_score/fos_score_prediction.py +218 -0
  20. libcom/fos_score/source/__init__.py +0 -0
  21. libcom/fos_score/source/config/__init__.py +1 -0
  22. libcom/fos_score/source/config/config.py +68 -0
  23. libcom/fos_score/source/config/config.sfosd.yaml +43 -0
  24. libcom/fos_score/source/config/config_rfosd.yaml +43 -0
  25. libcom/fos_score/source/loss_function.py +118 -0
  26. libcom/fos_score/source/networks.py +355 -0
  27. libcom/harmony_score/__init__.py +3 -0
  28. libcom/harmony_score/harmony_score_prediction.py +119 -0
  29. libcom/harmony_score/source/__init__.py +0 -0
  30. libcom/harmony_score/source/bargainnet.py +226 -0
  31. libcom/image_harmonization/__init__.py +4 -0
  32. libcom/image_harmonization/image_harmonization.py +193 -0
  33. libcom/image_harmonization/source/__init__.py +0 -0
  34. libcom/image_harmonization/source/functions.py +747 -0
  35. libcom/image_harmonization/source/pct_net.py +85 -0
  36. libcom/image_harmonization/source/src/__init__.py +0 -0
  37. libcom/image_harmonization/source/src/lbm/__init__.py +0 -0
  38. libcom/image_harmonization/source/src/lbm/config.py +141 -0
  39. libcom/image_harmonization/source/src/lbm/data/__init__.py +62 -0
  40. libcom/image_harmonization/source/src/lbm/data/datasets/__init__.py +9 -0
  41. libcom/image_harmonization/source/src/lbm/data/datasets/collation_fn.py +41 -0
  42. libcom/image_harmonization/source/src/lbm/data/datasets/dataset.py +243 -0
  43. libcom/image_harmonization/source/src/lbm/data/datasets/datasets_config.py +42 -0
  44. libcom/image_harmonization/source/src/lbm/data/filters/__init__.py +12 -0
  45. libcom/image_harmonization/source/src/lbm/data/filters/base.py +21 -0
  46. libcom/image_harmonization/source/src/lbm/data/filters/filter_wrapper.py +36 -0
  47. libcom/image_harmonization/source/src/lbm/data/filters/filters.py +33 -0
  48. libcom/image_harmonization/source/src/lbm/data/filters/filters_config.py +32 -0
  49. libcom/image_harmonization/source/src/lbm/data/mappers/__init__.py +19 -0
  50. libcom/image_harmonization/source/src/lbm/data/mappers/base.py +26 -0
  51. libcom/image_harmonization/source/src/lbm/data/mappers/mappers.py +135 -0
  52. libcom/image_harmonization/source/src/lbm/data/mappers/mappers_config.py +109 -0
  53. libcom/image_harmonization/source/src/lbm/data/mappers/mappers_wrapper.py +31 -0
  54. libcom/image_harmonization/source/src/lbm/inference/__init__.py +4 -0
  55. libcom/image_harmonization/source/src/lbm/inference/inference.py +70 -0
  56. libcom/image_harmonization/source/src/lbm/inference/utils.py +222 -0
  57. libcom/image_harmonization/source/src/lbm/models/__init__.py +0 -0
  58. libcom/image_harmonization/source/src/lbm/models/base/__init__.py +4 -0
  59. libcom/image_harmonization/source/src/lbm/models/base/base_model.py +66 -0
  60. libcom/image_harmonization/source/src/lbm/models/base/model_config.py +8 -0
  61. libcom/image_harmonization/source/src/lbm/models/embedders/__init__.py +4 -0
  62. libcom/image_harmonization/source/src/lbm/models/embedders/base/__init__.py +4 -0
  63. libcom/image_harmonization/source/src/lbm/models/embedders/base/base_conditioner.py +60 -0
  64. libcom/image_harmonization/source/src/lbm/models/embedders/base/base_conditioner_config.py +27 -0
  65. libcom/image_harmonization/source/src/lbm/models/embedders/conditioners_wrapper.py +114 -0
  66. libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/__init__.py +4 -0
  67. libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/latents_concat_embedder_config.py +31 -0
  68. libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/latents_concat_embedder_model.py +80 -0
  69. libcom/image_harmonization/source/src/lbm/models/lbm/__init__.py +4 -0
  70. libcom/image_harmonization/source/src/lbm/models/lbm/lbm_config.py +101 -0
  71. libcom/image_harmonization/source/src/lbm/models/lbm/lbm_model.py +359 -0
  72. libcom/image_harmonization/source/src/lbm/models/unets/__init__.py +14 -0
  73. libcom/image_harmonization/source/src/lbm/models/unets/unet.py +148 -0
  74. libcom/image_harmonization/source/src/lbm/models/utils.py +377 -0
  75. libcom/image_harmonization/source/src/lbm/models/vae/__init__.py +4 -0
  76. libcom/image_harmonization/source/src/lbm/models/vae/autoencoderKL.py +136 -0
  77. libcom/image_harmonization/source/src/lbm/models/vae/autoencoderKL_config.py +27 -0
  78. libcom/image_harmonization/source/src/lbm/trainer/__init__.py +85 -0
  79. libcom/image_harmonization/source/src/lbm/trainer/loggers.py +324 -0
  80. libcom/image_harmonization/source/src/lbm/trainer/trainer.py +199 -0
  81. libcom/image_harmonization/source/src/lbm/trainer/training_config.py +82 -0
  82. libcom/image_harmonization/source/src/lbm/trainer/utils.py +193 -0
  83. libcom/inharmonious_region_localization/__init__.py +4 -0
  84. libcom/inharmonious_region_localization/inharmonious_region_localization.py +97 -0
  85. libcom/inharmonious_region_localization/source/__init__.py +0 -0
  86. libcom/inharmonious_region_localization/source/blocks.py +368 -0
  87. libcom/inharmonious_region_localization/source/functions.py +679 -0
  88. libcom/inharmonious_region_localization/source/madis_net.py +79 -0
  89. libcom/kontext_blending_harmonization/__init__.py +3 -0
  90. libcom/kontext_blending_harmonization/kontext_blending_harmonization.py +148 -0
  91. libcom/kontext_blending_harmonization/source/__init__.py +0 -0
  92. libcom/kontext_blending_harmonization/source/pipeline.py +1135 -0
  93. libcom/naive_composition/__init__.py +3 -0
  94. libcom/naive_composition/generate_composite_image.py +116 -0
  95. libcom/opa_score/__init__.py +3 -0
  96. libcom/opa_score/opa_score_prediction.py +101 -0
  97. libcom/opa_score/source/__init__.py +3 -0
  98. libcom/opa_score/source/object_place_net.py +55 -0
  99. libcom/opa_score/source/oblect_place_config.py +66 -0
  100. libcom/opa_score/source/resnet_4ch.py +213 -0
  101. libcom/os_insert/__init__.py +3 -0
  102. libcom/os_insert/diffusers_osinsert/__init__.py +78 -0
  103. libcom/os_insert/diffusers_osinsert/_patched_pipeline_flux_fill.py +878 -0
  104. libcom/os_insert/os_insert.py +534 -0
  105. libcom/os_insert/source/__init__.py +1 -0
  106. libcom/os_insert/source/ia_utils.py +183 -0
  107. libcom/os_insert/source/insertanything_infer.py +529 -0
  108. libcom/os_insert/source/ldm/__init__.py +0 -0
  109. libcom/os_insert/source/ldm/data/__init__.py +0 -0
  110. libcom/os_insert/source/ldm/data/base.py +23 -0
  111. libcom/os_insert/source/ldm/data/datamodule.py +128 -0
  112. libcom/os_insert/source/ldm/data/imagenet.py +394 -0
  113. libcom/os_insert/source/ldm/data/lsun.py +92 -0
  114. libcom/os_insert/source/ldm/data/open_images.py +728 -0
  115. libcom/os_insert/source/ldm/lr_scheduler.py +98 -0
  116. libcom/os_insert/source/ldm/models/__init__.py +0 -0
  117. libcom/os_insert/source/ldm/models/autoencoder.py +443 -0
  118. libcom/os_insert/source/ldm/models/diffusion/__init__.py +0 -0
  119. libcom/os_insert/source/ldm/models/diffusion/classifier.py +267 -0
  120. libcom/os_insert/source/ldm/models/diffusion/ddim.py +263 -0
  121. libcom/os_insert/source/ldm/models/diffusion/ddpm.py +1417 -0
  122. libcom/os_insert/source/ldm/models/diffusion/plms.py +252 -0
  123. libcom/os_insert/source/ldm/modules/__init__.py +0 -0
  124. libcom/os_insert/source/ldm/modules/attention.py +261 -0
  125. libcom/os_insert/source/ldm/modules/diffusionmodules/__init__.py +0 -0
  126. libcom/os_insert/source/ldm/modules/diffusionmodules/model.py +835 -0
  127. libcom/os_insert/source/ldm/modules/diffusionmodules/openaimodel.py +1108 -0
  128. libcom/os_insert/source/ldm/modules/diffusionmodules/util.py +267 -0
  129. libcom/os_insert/source/ldm/modules/distributions/__init__.py +0 -0
  130. libcom/os_insert/source/ldm/modules/distributions/distributions.py +92 -0
  131. libcom/os_insert/source/ldm/modules/ema.py +76 -0
  132. libcom/os_insert/source/ldm/modules/encoders/__init__.py +0 -0
  133. libcom/os_insert/source/ldm/modules/encoders/modules.py +190 -0
  134. libcom/os_insert/source/ldm/modules/encoders/xf.py +130 -0
  135. libcom/os_insert/source/ldm/modules/losses/__init__.py +1 -0
  136. libcom/os_insert/source/ldm/modules/losses/contperceptual.py +111 -0
  137. libcom/os_insert/source/ldm/modules/losses/vqperceptual.py +167 -0
  138. libcom/os_insert/source/ldm/modules/x_transformer.py +641 -0
  139. libcom/os_insert/source/ldm/util.py +203 -0
  140. libcom/os_insert/source/objectstitch_data.py +66 -0
  141. libcom/os_insert/source/objectstitch_infer.py +618 -0
  142. libcom/os_insert/source/sam_on_objectstitch.py +253 -0
  143. libcom/os_insert/source/utils.py +41 -0
  144. libcom/os_insert/source/v1.yaml +108 -0
  145. libcom/painterly_image_harmonization/__init__.py +4 -0
  146. libcom/painterly_image_harmonization/painterly_image_harmonization.py +215 -0
  147. libcom/painterly_image_harmonization/source/PHDNet/__init__.py +0 -0
  148. libcom/painterly_image_harmonization/source/PHDNet/networks.py +202 -0
  149. libcom/painterly_image_harmonization/source/PHDNet/phdnet.py +47 -0
  150. libcom/painterly_image_harmonization/source/PHDNet/phdnet_config.py +18 -0
  151. libcom/painterly_image_harmonization/source/PHDiffusion/__init__.py +0 -0
  152. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/__init__.py +0 -0
  153. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/inference_base.py +300 -0
  154. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/lr_scheduler.py +98 -0
  155. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/__init__.py +0 -0
  156. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/autoencoder.py +211 -0
  157. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/__init__.py +3 -0
  158. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/ddim.py +336 -0
  159. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/ddpm.py +1400 -0
  160. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
  161. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/dpm_solver/dpm_solver.py +1217 -0
  162. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/dpm_solver/sampler.py +87 -0
  163. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/plms.py +243 -0
  164. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/scheduling_pndm.py +518 -0
  165. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/__init__.py +0 -0
  166. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/attention.py +378 -0
  167. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/__init__.py +0 -0
  168. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/model.py +857 -0
  169. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/openaimodel.py +823 -0
  170. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/util.py +270 -0
  171. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/distributions/__init__.py +0 -0
  172. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/distributions/distributions.py +92 -0
  173. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/ema.py +80 -0
  174. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/__init__.py +0 -0
  175. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/adapter.py +213 -0
  176. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/modules.py +82 -0
  177. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/__init__.py +1 -0
  178. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/api.py +269 -0
  179. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/__init__.py +0 -0
  180. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/api.py +175 -0
  181. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/__init__.py +0 -0
  182. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/base_model.py +16 -0
  183. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/blocks.py +342 -0
  184. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/dpt_depth.py +109 -0
  185. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/midas_net.py +76 -0
  186. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/midas_net_custom.py +128 -0
  187. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/transforms.py +234 -0
  188. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/vit.py +491 -0
  189. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/utils.py +189 -0
  190. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/model_edge.py +653 -0
  191. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/__init__.py +0 -0
  192. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/api.py +35 -0
  193. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/body.py +209 -0
  194. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/hand.py +75 -0
  195. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/model.py +178 -0
  196. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/util.py +203 -0
  197. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/utils.py +72 -0
  198. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/fusion/__init__.py +0 -0
  199. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/fusion/dual_encoder_fusion.py +93 -0
  200. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/__init__.py +2 -0
  201. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/bsrgan.py +730 -0
  202. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/bsrgan_light.py +651 -0
  203. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/utils_image.py +916 -0
  204. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/util.py +200 -0
  205. libcom/painterly_image_harmonization/source/PHDiffusion/stable_diffusion.yaml +75 -0
  206. libcom/painterly_image_harmonization/source/__init__.py +0 -0
  207. libcom/reflection_generation/__init__.py +3 -0
  208. libcom/reflection_generation/reflection_generation.py +174 -0
  209. libcom/reflection_generation/source/PostProcessModel.py +513 -0
  210. libcom/reflection_generation/source/__init__.py +0 -0
  211. libcom/reflection_generation/source/cldm/__init__.py +0 -0
  212. libcom/reflection_generation/source/cldm/base_network.py +166 -0
  213. libcom/reflection_generation/source/cldm/cldm.py +509 -0
  214. libcom/reflection_generation/source/cldm/ddim_hacked.py +317 -0
  215. libcom/reflection_generation/source/cldm/hack.py +111 -0
  216. libcom/reflection_generation/source/cldm/logger.py +127 -0
  217. libcom/reflection_generation/source/cldm/model.py +26 -0
  218. libcom/reflection_generation/source/cldm_v15.yaml +81 -0
  219. libcom/reflection_generation/source/cluster_feature.py +109 -0
  220. libcom/reflection_generation/source/ldm/__init__.py +0 -0
  221. libcom/reflection_generation/source/ldm/models/__init__.py +0 -0
  222. libcom/reflection_generation/source/ldm/models/autoencoder.py +219 -0
  223. libcom/reflection_generation/source/ldm/models/diffusion/__init__.py +0 -0
  224. libcom/reflection_generation/source/ldm/models/diffusion/attention_processor.py +188 -0
  225. libcom/reflection_generation/source/ldm/models/diffusion/ddim.py +340 -0
  226. libcom/reflection_generation/source/ldm/models/diffusion/ddpm.py +1892 -0
  227. libcom/reflection_generation/source/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
  228. libcom/reflection_generation/source/ldm/models/diffusion/dpm_solver/dpm_solver.py +1154 -0
  229. libcom/reflection_generation/source/ldm/models/diffusion/dpm_solver/sampler.py +87 -0
  230. libcom/reflection_generation/source/ldm/models/diffusion/ip_adapter.py +57 -0
  231. libcom/reflection_generation/source/ldm/models/diffusion/plms.py +244 -0
  232. libcom/reflection_generation/source/ldm/models/diffusion/sampling_util.py +22 -0
  233. libcom/reflection_generation/source/ldm/modules/__init__.py +0 -0
  234. libcom/reflection_generation/source/ldm/modules/attention.py +341 -0
  235. libcom/reflection_generation/source/ldm/modules/diffusionmodules/__init__.py +0 -0
  236. libcom/reflection_generation/source/ldm/modules/diffusionmodules/model.py +852 -0
  237. libcom/reflection_generation/source/ldm/modules/diffusionmodules/openaimodel.py +786 -0
  238. libcom/reflection_generation/source/ldm/modules/diffusionmodules/upscaling.py +81 -0
  239. libcom/reflection_generation/source/ldm/modules/diffusionmodules/util.py +280 -0
  240. libcom/reflection_generation/source/ldm/modules/distributions/__init__.py +0 -0
  241. libcom/reflection_generation/source/ldm/modules/distributions/distributions.py +92 -0
  242. libcom/reflection_generation/source/ldm/modules/ema.py +80 -0
  243. libcom/reflection_generation/source/ldm/modules/encoders/__init__.py +0 -0
  244. libcom/reflection_generation/source/ldm/modules/encoders/modules.py +214 -0
  245. libcom/reflection_generation/source/ldm/modules/image_degradation/__init__.py +0 -0
  246. libcom/reflection_generation/source/ldm/modules/image_degradation/bsrgan.py +730 -0
  247. libcom/reflection_generation/source/ldm/modules/image_degradation/bsrgan_light.py +651 -0
  248. libcom/reflection_generation/source/ldm/modules/image_degradation/utils_image.py +916 -0
  249. libcom/reflection_generation/source/ldm/util.py +198 -0
  250. libcom/reflection_generation/source/local_clip.py +749 -0
  251. libcom/reflection_generation/source/tool_add_control.py +49 -0
  252. libcom/shadow_generation/__init__.py +3 -0
  253. libcom/shadow_generation/shadow_generation.py +754 -0
  254. libcom/shadow_generation/source/PostProcessModel.py +511 -0
  255. libcom/shadow_generation/source/__init__.py +0 -0
  256. libcom/shadow_generation/source/attention_processor.py +234 -0
  257. libcom/shadow_generation/source/cldm/__init__.py +0 -0
  258. libcom/shadow_generation/source/cldm/base_network.py +166 -0
  259. libcom/shadow_generation/source/cldm/cldm.py +524 -0
  260. libcom/shadow_generation/source/cldm/model.py +28 -0
  261. libcom/shadow_generation/source/ldm/__init__.py +0 -0
  262. libcom/shadow_generation/source/ldm/models/__init__.py +0 -0
  263. libcom/shadow_generation/source/ldm/models/autoencoder.py +219 -0
  264. libcom/shadow_generation/source/ldm/models/diffusion/__init__.py +0 -0
  265. libcom/shadow_generation/source/ldm/models/diffusion/attention_processor.py +188 -0
  266. libcom/shadow_generation/source/ldm/models/diffusion/ddim.py +340 -0
  267. libcom/shadow_generation/source/ldm/models/diffusion/ddpm.py +1893 -0
  268. libcom/shadow_generation/source/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
  269. libcom/shadow_generation/source/ldm/models/diffusion/dpm_solver/dpm_solver.py +1154 -0
  270. libcom/shadow_generation/source/ldm/models/diffusion/dpm_solver/sampler.py +87 -0
  271. libcom/shadow_generation/source/ldm/models/diffusion/ip_adapter.py +57 -0
  272. libcom/shadow_generation/source/ldm/models/diffusion/plms.py +244 -0
  273. libcom/shadow_generation/source/ldm/models/diffusion/sampling_util.py +22 -0
  274. libcom/shadow_generation/source/ldm/modules/__init__.py +0 -0
  275. libcom/shadow_generation/source/ldm/modules/attention.py +341 -0
  276. libcom/shadow_generation/source/ldm/modules/diffusionmodules/__init__.py +0 -0
  277. libcom/shadow_generation/source/ldm/modules/diffusionmodules/model.py +852 -0
  278. libcom/shadow_generation/source/ldm/modules/diffusionmodules/openaimodel.py +786 -0
  279. libcom/shadow_generation/source/ldm/modules/diffusionmodules/upscaling.py +81 -0
  280. libcom/shadow_generation/source/ldm/modules/diffusionmodules/util.py +270 -0
  281. libcom/shadow_generation/source/ldm/modules/distributions/__init__.py +0 -0
  282. libcom/shadow_generation/source/ldm/modules/distributions/distributions.py +92 -0
  283. libcom/shadow_generation/source/ldm/modules/ema.py +80 -0
  284. libcom/shadow_generation/source/ldm/modules/encoders/__init__.py +0 -0
  285. libcom/shadow_generation/source/ldm/modules/encoders/modules.py +214 -0
  286. libcom/shadow_generation/source/ldm/modules/image_degradation/__init__.py +2 -0
  287. libcom/shadow_generation/source/ldm/modules/image_degradation/bsrgan.py +730 -0
  288. libcom/shadow_generation/source/ldm/modules/image_degradation/bsrgan_light.py +651 -0
  289. libcom/shadow_generation/source/ldm/modules/image_degradation/utils_image.py +916 -0
  290. libcom/shadow_generation/source/ldm/modules/midas/__init__.py +0 -0
  291. libcom/shadow_generation/source/ldm/modules/midas/api.py +170 -0
  292. libcom/shadow_generation/source/ldm/modules/midas/midas/__init__.py +0 -0
  293. libcom/shadow_generation/source/ldm/modules/midas/midas/base_model.py +16 -0
  294. libcom/shadow_generation/source/ldm/modules/midas/midas/blocks.py +342 -0
  295. libcom/shadow_generation/source/ldm/modules/midas/midas/dpt_depth.py +109 -0
  296. libcom/shadow_generation/source/ldm/modules/midas/midas/midas_net.py +76 -0
  297. libcom/shadow_generation/source/ldm/modules/midas/midas/midas_net_custom.py +128 -0
  298. libcom/shadow_generation/source/ldm/modules/midas/midas/transforms.py +234 -0
  299. libcom/shadow_generation/source/ldm/modules/midas/midas/vit.py +491 -0
  300. libcom/shadow_generation/source/ldm/modules/midas/utils.py +189 -0
  301. libcom/shadow_generation/source/ldm/util.py +199 -0
  302. libcom/utils/__init__.py +0 -0
  303. libcom/utils/environment.py +12 -0
  304. libcom/utils/model_download.py +67 -0
  305. libcom/utils/process_image.py +108 -0
  306. libcom-0.1.8.dist-info/METADATA +157 -0
  307. libcom-0.1.8.dist-info/RECORD +310 -0
  308. libcom-0.1.8.dist-info/WHEEL +5 -0
  309. libcom-0.1.8.dist-info/licenses/LICENSE +201 -0
  310. libcom-0.1.8.dist-info/top_level.txt +1 -0
@@ -0,0 +1,749 @@
1
+
2
+ from functools import partial
3
+ from torch import nn, einsum
4
+ import torch.nn.functional as F
5
+ import torch
6
+ from inspect import isfunction
7
+ from einops import rearrange, repeat, reduce
8
+ from collections import namedtuple
9
+ from transformers import CLIPVisionModel
10
+ import math
11
+
12
+
13
+ DEFAULT_DIM_HEAD = 64
14
+
15
+
16
+ Intermediates = namedtuple('Intermediates', [
17
+ 'pre_softmax_attn',
18
+ 'post_softmax_attn'
19
+ ])
20
+
21
+
22
+ LayerIntermediates = namedtuple('Intermediates', [
23
+ 'hiddens',
24
+ 'attn_intermediates'
25
+ ])
26
+
27
+
28
+ def exists(val):
29
+ return val is not None
30
+
31
+
32
+ def always(val):
33
+ def inner(*args, **kwargs):
34
+ return val
35
+ return inner
36
+
37
+
38
+ def not_equals(val):
39
+ def inner(x):
40
+ return x != val
41
+ return inner
42
+
43
+
44
+ def default(val, d):
45
+ if exists(val):
46
+ return val
47
+ return d() if isfunction(d) else d
48
+
49
+
50
+ def string_begins_with(prefix, str):
51
+ return str.startswith(prefix)
52
+
53
+
54
+ def group_dict_by_key(cond, d):
55
+ return_val = [dict(), dict()]
56
+ for key in d.keys():
57
+ match = bool(cond(key))
58
+ ind = int(not match)
59
+ return_val[ind][key] = d[key]
60
+ return (*return_val,)
61
+
62
+
63
+ class Scale(nn.Module):
64
+ def __init__(self, value, fn):
65
+ super().__init__()
66
+ self.value = value
67
+ self.fn = fn
68
+
69
+ def forward(self, x, **kwargs):
70
+ x, *rest = self.fn(x, **kwargs)
71
+ return (x * self.value, *rest)
72
+
73
+
74
+ class Residual(nn.Module):
75
+ def forward(self, x, residual):
76
+ return x + residual
77
+
78
+
79
+ class AbsolutePositionalEmbedding(nn.Module):
80
+ def __init__(self, dim, max_seq_len):
81
+ super().__init__()
82
+ self.emb = nn.Embedding(max_seq_len, dim)
83
+ self.init_()
84
+
85
+ def init_(self):
86
+ nn.init.normal_(self.emb.weight, std=0.02)
87
+
88
+ def forward(self, x):
89
+ n = torch.arange(x.shape[1], device=x.device)
90
+ return self.emb(n)[None, :, :]
91
+
92
+
93
+ class GRUGating(nn.Module):
94
+ def __init__(self, dim):
95
+ super().__init__()
96
+ self.gru = nn.GRUCell(dim, dim)
97
+
98
+ def forward(self, x, residual):
99
+ gated_output = self.gru(
100
+ rearrange(x, 'b n d -> (b n) d'),
101
+ rearrange(residual, 'b n d -> (b n) d')
102
+ )
103
+
104
+ return gated_output.reshape_as(x)
105
+
106
+
107
+ class Rezero(nn.Module):
108
+ def __init__(self, fn):
109
+ super().__init__()
110
+ self.fn = fn
111
+ self.g = nn.Parameter(torch.zeros(1))
112
+
113
+ def forward(self, x, **kwargs):
114
+ x, *rest = self.fn(x, **kwargs)
115
+ return (x * self.g, *rest)
116
+
117
+
118
+ class ScaleNorm(nn.Module):
119
+ def __init__(self, dim, eps=1e-5):
120
+ super().__init__()
121
+ self.scale = dim ** -0.5
122
+ self.eps = eps
123
+ self.g = nn.Parameter(torch.ones(1))
124
+
125
+ def forward(self, x):
126
+ norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
127
+ return x / norm.clamp(min=self.eps) * self.g
128
+
129
+
130
+ def groupby_prefix_and_trim(prefix, d):
131
+ kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, prefix), d)
132
+ kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items())))
133
+ return kwargs_without_prefix, kwargs
134
+
135
+
136
+ class RMSNorm(nn.Module):
137
+ def __init__(self, dim, eps=1e-8):
138
+ super().__init__()
139
+ self.scale = dim ** -0.5
140
+ self.eps = eps
141
+ self.g = nn.Parameter(torch.ones(dim))
142
+
143
+ def forward(self, x):
144
+ norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
145
+ return x / norm.clamp(min=self.eps) * self.g
146
+
147
+
148
+ class FixedPositionalEmbedding(nn.Module):
149
+ def __init__(self, dim):
150
+ super().__init__()
151
+ inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
152
+ self.register_buffer('inv_freq', inv_freq)
153
+
154
+ def forward(self, x, seq_dim=1, offset=0):
155
+ t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) + offset
156
+ sinusoid_inp = torch.einsum('i , j -> i j', t, self.inv_freq)
157
+ emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1)
158
+ return emb[None, :, :]
159
+
160
+
161
+ def max_neg_value(tensor):
162
+ return -torch.finfo(tensor.dtype).max
163
+
164
+
165
+ class TransformerWrapper(nn.Module):
166
+ def __init__(
167
+ self,
168
+ *,
169
+ num_tokens,
170
+ max_seq_len,
171
+ attn_layers,
172
+ emb_dim=None,
173
+ max_mem_len=0.,
174
+ emb_dropout=0.,
175
+ num_memory_tokens=None,
176
+ tie_embedding=False,
177
+ use_pos_emb=True
178
+ ):
179
+ super().__init__()
180
+ assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
181
+
182
+ dim = attn_layers.dim
183
+ emb_dim = default(emb_dim, dim)
184
+
185
+ self.max_seq_len = max_seq_len
186
+ self.max_mem_len = max_mem_len
187
+ self.num_tokens = num_tokens
188
+
189
+ self.token_emb = nn.Embedding(num_tokens, emb_dim)
190
+ self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len) if (
191
+ use_pos_emb and not attn_layers.has_pos_emb) else always(0)
192
+ self.emb_dropout = nn.Dropout(emb_dropout)
193
+
194
+ self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
195
+ self.attn_layers = attn_layers
196
+ self.norm = nn.LayerNorm(dim)
197
+
198
+ self.init_()
199
+
200
+ self.to_logits = nn.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t()
201
+
202
+ # memory tokens (like [cls]) from Memory Transformers paper
203
+ num_memory_tokens = default(num_memory_tokens, 0)
204
+ self.num_memory_tokens = num_memory_tokens
205
+ if num_memory_tokens > 0:
206
+ self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
207
+
208
+ # let funnel encoder know number of memory tokens, if specified
209
+ if hasattr(attn_layers, 'num_memory_tokens'):
210
+ attn_layers.num_memory_tokens = num_memory_tokens
211
+
212
+
213
+ def not_equals(val):
214
+ def inner(x):
215
+ return x != val
216
+ return inner
217
+
218
+
219
+ def equals(val):
220
+ def inner(x):
221
+ return x == val
222
+ return inner
223
+
224
+
225
+ class Attention(nn.Module):
226
+ def __init__(
227
+ self,
228
+ dim,
229
+ dim_head=DEFAULT_DIM_HEAD,
230
+ heads=8,
231
+ causal=False,
232
+ mask=None,
233
+ talking_heads=False,
234
+ sparse_topk=None,
235
+ use_entmax15=False,
236
+ num_mem_kv=0,
237
+ dropout=0.,
238
+ on_attn=False
239
+ ):
240
+ super().__init__()
241
+ if use_entmax15:
242
+ raise NotImplementedError("Check out entmax activation instead of softmax activation!")
243
+ self.scale = dim_head ** -0.5
244
+ self.heads = heads
245
+ self.causal = causal
246
+ self.mask = mask
247
+
248
+ inner_dim = dim_head * heads
249
+
250
+ self.to_q = nn.Linear(dim, inner_dim, bias=False)
251
+ self.to_k = nn.Linear(dim, inner_dim, bias=False)
252
+ self.to_v = nn.Linear(dim, inner_dim, bias=False)
253
+ self.dropout = nn.Dropout(dropout)
254
+
255
+ # talking heads
256
+ self.talking_heads = talking_heads
257
+ if talking_heads:
258
+ self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads))
259
+ self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads))
260
+
261
+ # explicit topk sparse attention
262
+ self.sparse_topk = sparse_topk
263
+
264
+ # entmax
265
+ #self.attn_fn = entmax15 if use_entmax15 else F.softmax
266
+ self.attn_fn = F.softmax
267
+
268
+ # add memory key / values
269
+ self.num_mem_kv = num_mem_kv
270
+ if num_mem_kv > 0:
271
+ self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
272
+ self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
273
+
274
+ # attention on attention
275
+ self.attn_on_attn = on_attn
276
+ self.to_out = nn.Sequential(nn.Linear(inner_dim, dim * 2), nn.GLU()) if on_attn else nn.Linear(inner_dim, dim)
277
+
278
+ def forward(
279
+ self,
280
+ x,
281
+ context=None,
282
+ mask=None,
283
+ context_mask=None,
284
+ rel_pos=None,
285
+ sinusoidal_emb=None,
286
+ prev_attn=None,
287
+ mem=None
288
+ ):
289
+ b, n, _, h, talking_heads, device = *x.shape, self.heads, self.talking_heads, x.device
290
+ kv_input = default(context, x)
291
+
292
+ q_input = x
293
+ k_input = kv_input
294
+ v_input = kv_input
295
+
296
+ if exists(mem):
297
+ k_input = torch.cat((mem, k_input), dim=-2)
298
+ v_input = torch.cat((mem, v_input), dim=-2)
299
+
300
+ if exists(sinusoidal_emb):
301
+ # in shortformer, the query would start at a position offset depending on the past cached memory
302
+ offset = k_input.shape[-2] - q_input.shape[-2]
303
+ q_input = q_input + sinusoidal_emb(q_input, offset=offset)
304
+ k_input = k_input + sinusoidal_emb(k_input)
305
+
306
+ q = self.to_q(q_input)
307
+ k = self.to_k(k_input)
308
+ v = self.to_v(v_input)
309
+
310
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
311
+
312
+ input_mask = None
313
+ if any(map(exists, (mask, context_mask))):
314
+ q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool())
315
+ k_mask = q_mask if not exists(context) else context_mask
316
+ k_mask = default(k_mask, lambda: torch.ones((b, k.shape[-2]), device=device).bool())
317
+ q_mask = rearrange(q_mask, 'b i -> b () i ()')
318
+ k_mask = rearrange(k_mask, 'b j -> b () () j')
319
+ input_mask = q_mask * k_mask
320
+
321
+ if self.num_mem_kv > 0:
322
+ mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b=b), (self.mem_k, self.mem_v))
323
+ k = torch.cat((mem_k, k), dim=-2)
324
+ v = torch.cat((mem_v, v), dim=-2)
325
+ if exists(input_mask):
326
+ input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True)
327
+
328
+ dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
329
+ mask_value = max_neg_value(dots)
330
+
331
+ if exists(prev_attn):
332
+ dots = dots + prev_attn
333
+
334
+ pre_softmax_attn = dots
335
+
336
+ if talking_heads:
337
+ dots = einsum('b h i j, h k -> b k i j', dots, self.pre_softmax_proj).contiguous()
338
+
339
+ if exists(rel_pos):
340
+ dots = rel_pos(dots)
341
+
342
+ if exists(input_mask):
343
+ dots.masked_fill_(~input_mask, mask_value)
344
+ del input_mask
345
+
346
+ if self.causal:
347
+ i, j = dots.shape[-2:]
348
+ r = torch.arange(i, device=device)
349
+ mask = rearrange(r, 'i -> () () i ()') < rearrange(r, 'j -> () () () j')
350
+ mask = F.pad(mask, (j - i, 0), value=False)
351
+ dots.masked_fill_(mask, mask_value)
352
+ del mask
353
+
354
+ if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]:
355
+ top, _ = dots.topk(self.sparse_topk, dim=-1)
356
+ vk = top[..., -1].unsqueeze(-1).expand_as(dots)
357
+ mask = dots < vk
358
+ dots.masked_fill_(mask, mask_value)
359
+ del mask
360
+
361
+ attn = self.attn_fn(dots, dim=-1)
362
+ post_softmax_attn = attn
363
+
364
+ attn = self.dropout(attn)
365
+
366
+ if talking_heads:
367
+ attn = einsum('b h i j, h k -> b k i j', attn, self.post_softmax_proj).contiguous()
368
+
369
+ out = einsum('b h i j, b h j d -> b h i d', attn, v)
370
+ out = rearrange(out, 'b h n d -> b n (h d)')
371
+
372
+ intermediates = Intermediates(
373
+ pre_softmax_attn=pre_softmax_attn,
374
+ post_softmax_attn=post_softmax_attn
375
+ )
376
+
377
+ return self.to_out(out), intermediates
378
+
379
+
380
+ class GEGLU(nn.Module):
381
+ def __init__(self, dim_in, dim_out):
382
+ super().__init__()
383
+ self.proj = nn.Linear(dim_in, dim_out * 2)
384
+
385
+ def forward(self, x):
386
+ x, gate = self.proj(x).chunk(2, dim=-1)
387
+ return x * F.gelu(gate)
388
+
389
+
390
+ class FeedForward(nn.Module):
391
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
392
+ super().__init__()
393
+ inner_dim = int(dim * mult)
394
+ dim_out = default(dim_out, dim)
395
+ project_in = nn.Sequential(
396
+ nn.Linear(dim, inner_dim),
397
+ nn.GELU()
398
+ ) if not glu else GEGLU(dim, inner_dim)
399
+
400
+ self.net = nn.Sequential(
401
+ project_in,
402
+ nn.Dropout(dropout),
403
+ nn.Linear(inner_dim, dim_out)
404
+ )
405
+
406
+ def forward(self, x):
407
+ return self.net(x)
408
+
409
+
410
+ class AttentionLayers(nn.Module):
411
+ def __init__(
412
+ self,
413
+ dim,
414
+ depth,
415
+ heads=8,
416
+ causal=False,
417
+ cross_attend=False,
418
+ only_cross=False,
419
+ use_scalenorm=False,
420
+ use_rmsnorm=False,
421
+ use_rezero=False,
422
+ rel_pos_num_buckets=32,
423
+ rel_pos_max_distance=128,
424
+ position_infused_attn=False,
425
+ custom_layers=None,
426
+ sandwich_coef=None,
427
+ par_ratio=None,
428
+ residual_attn=False,
429
+ cross_residual_attn=False,
430
+ macaron=False,
431
+ pre_norm=True,
432
+ gate_residual=False,
433
+ **kwargs
434
+ ):
435
+ super().__init__()
436
+ ff_kwargs, kwargs = groupby_prefix_and_trim('ff_', kwargs)
437
+ attn_kwargs, _ = groupby_prefix_and_trim('attn_', kwargs)
438
+
439
+ dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
440
+
441
+ self.dim = dim
442
+ self.depth = depth
443
+ self.layers = nn.ModuleList([])
444
+
445
+ self.has_pos_emb = position_infused_attn
446
+ self.pia_pos_emb = FixedPositionalEmbedding(dim) if position_infused_attn else None
447
+ self.rotary_pos_emb = always(None)
448
+
449
+ assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance'
450
+ self.rel_pos = None
451
+
452
+ self.pre_norm = pre_norm
453
+
454
+ self.residual_attn = residual_attn
455
+ self.cross_residual_attn = cross_residual_attn
456
+
457
+ norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm
458
+ norm_class = RMSNorm if use_rmsnorm else norm_class
459
+ norm_fn = partial(norm_class, dim)
460
+
461
+ norm_fn = nn.Identity if use_rezero else norm_fn
462
+ branch_fn = Rezero if use_rezero else None
463
+
464
+ if cross_attend and not only_cross:
465
+ default_block = ('a', 'c', 'f')
466
+ elif cross_attend and only_cross:
467
+ default_block = ('c', 'f')
468
+ else:
469
+ default_block = ('a', 'f')
470
+
471
+ if macaron:
472
+ default_block = ('f',) + default_block
473
+
474
+ if exists(custom_layers):
475
+ layer_types = custom_layers
476
+ elif exists(par_ratio):
477
+ par_depth = depth * len(default_block)
478
+ assert 1 < par_ratio <= par_depth, 'par ratio out of range'
479
+ default_block = tuple(filter(not_equals('f'), default_block))
480
+ par_attn = par_depth // par_ratio
481
+ depth_cut = par_depth * 2 // 3 # 2 / 3 attention layer cutoff suggested by PAR paper
482
+ par_width = (depth_cut + depth_cut // par_attn) // par_attn
483
+ assert len(default_block) <= par_width, 'default block is too large for par_ratio'
484
+ par_block = default_block + ('f',) * (par_width - len(default_block))
485
+ par_head = par_block * par_attn
486
+ layer_types = par_head + ('f',) * (par_depth - len(par_head))
487
+ elif exists(sandwich_coef):
488
+ assert sandwich_coef > 0 and sandwich_coef <= depth, 'sandwich coefficient should be less than the depth'
489
+ layer_types = ('a',) * sandwich_coef + default_block * (depth - sandwich_coef) + ('f',) * sandwich_coef
490
+ else:
491
+ layer_types = default_block * depth
492
+
493
+ self.layer_types = layer_types
494
+ self.num_attn_layers = len(list(filter(equals('a'), layer_types)))
495
+
496
+ for layer_type in self.layer_types:
497
+ if layer_type == 'a':
498
+ layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs)
499
+ elif layer_type == 'c':
500
+ layer = Attention(dim, heads=heads, **attn_kwargs)
501
+ elif layer_type == 'f':
502
+ layer = FeedForward(dim, **ff_kwargs)
503
+ layer = layer if not macaron else Scale(0.5, layer)
504
+ else:
505
+ raise Exception(f'invalid layer type {layer_type}')
506
+
507
+ if isinstance(layer, Attention) and exists(branch_fn):
508
+ layer = branch_fn(layer)
509
+
510
+ if gate_residual:
511
+ residual_fn = GRUGating(dim)
512
+ else:
513
+ residual_fn = Residual()
514
+
515
+ self.layers.append(nn.ModuleList([
516
+ norm_fn(),
517
+ layer,
518
+ residual_fn
519
+ ]))
520
+
521
+ def forward(
522
+ self,
523
+ x,
524
+ context=None,
525
+ mask=None,
526
+ context_mask=None,
527
+ mems=None,
528
+ return_hiddens=False
529
+ ):
530
+ hiddens = []
531
+ intermediates = []
532
+ prev_attn = None
533
+ prev_cross_attn = None
534
+
535
+ mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers
536
+
537
+ for ind, (layer_type, (norm, block, residual_fn)) in enumerate(zip(self.layer_types, self.layers)):
538
+ is_last = ind == (len(self.layers) - 1)
539
+
540
+ if layer_type == 'a':
541
+ hiddens.append(x)
542
+ layer_mem = mems.pop(0)
543
+
544
+ residual = x
545
+
546
+ if self.pre_norm:
547
+ x = norm(x)
548
+
549
+ if layer_type == 'a':
550
+ out, inter = block(x, mask=mask, sinusoidal_emb=self.pia_pos_emb, rel_pos=self.rel_pos,
551
+ prev_attn=prev_attn, mem=layer_mem)
552
+ elif layer_type == 'c':
553
+ out, inter = block(x, context=context, mask=mask, context_mask=context_mask, prev_attn=prev_cross_attn)
554
+ elif layer_type == 'f':
555
+ out = block(x)
556
+
557
+ x = residual_fn(out, residual)
558
+
559
+ if layer_type in ('a', 'c'):
560
+ intermediates.append(inter)
561
+
562
+ if layer_type == 'a' and self.residual_attn:
563
+ prev_attn = inter.pre_softmax_attn
564
+ elif layer_type == 'c' and self.cross_residual_attn:
565
+ prev_cross_attn = inter.pre_softmax_attn
566
+
567
+ if not self.pre_norm and not is_last:
568
+ x = norm(x)
569
+
570
+ if return_hiddens:
571
+ intermediates = LayerIntermediates(
572
+ hiddens=hiddens,
573
+ attn_intermediates=intermediates
574
+ )
575
+
576
+ return x, intermediates
577
+
578
+ return x
579
+
580
+
581
+ class Encoder(AttentionLayers):
582
+ def __init__(self, **kwargs):
583
+ assert 'causal' not in kwargs, 'cannot set causality on encoder'
584
+ super().__init__(causal=False, **kwargs)
585
+
586
+
587
+ class AbstractEncoder(nn.Module):
588
+ def __init__(self):
589
+ super().__init__()
590
+
591
+ def encode(self, *args, **kwargs):
592
+ raise NotImplementedError
593
+
594
+
595
+ class MultiheadAttention(nn.Module):
596
+ def __init__(self, n_ctx, width, heads):
597
+ super().__init__()
598
+ self.n_ctx = n_ctx
599
+ self.width = width
600
+ self.heads = heads
601
+ self.c_qkv = nn.Linear(width, width * 3)
602
+ self.c_proj = nn.Linear(width, width)
603
+ self.attention = QKVMultiheadAttention(heads, n_ctx)
604
+
605
+ def forward(self, x):
606
+ x = self.c_qkv(x)
607
+ x = self.attention(x)
608
+ x = self.c_proj(x)
609
+ return x
610
+
611
+
612
+ class MLP(nn.Module):
613
+ def __init__(self, width):
614
+ super().__init__()
615
+ self.width = width
616
+ self.c_fc = nn.Linear(width, width * 4)
617
+ self.c_proj = nn.Linear(width * 4, width)
618
+ self.gelu = nn.GELU()
619
+
620
+ def forward(self, x):
621
+ return self.c_proj(self.gelu(self.c_fc(x)))
622
+
623
+
624
+ class QKVMultiheadAttention(nn.Module):
625
+ def __init__(self, n_heads: int, n_ctx: int):
626
+ super().__init__()
627
+ self.n_heads = n_heads
628
+ self.n_ctx = n_ctx
629
+
630
+ def forward(self, qkv):
631
+ bs, n_ctx, width = qkv.shape
632
+ attn_ch = width // self.n_heads // 3
633
+ scale = 1 / math.sqrt(math.sqrt(attn_ch))
634
+ qkv = qkv.view(bs, n_ctx, self.n_heads, -1)
635
+ q, k, v = torch.split(qkv, attn_ch, dim=-1)
636
+ weight = torch.einsum(
637
+ "bthc,bshc->bhts", q * scale, k * scale
638
+ ) # More stable with f16 than dividing afterwards
639
+ wdtype = weight.dtype
640
+ weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
641
+ return torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
642
+
643
+
644
+ class ResidualAttentionBlock(nn.Module):
645
+ def __init__(
646
+ self,
647
+ n_ctx: int,
648
+ width: int,
649
+ heads: int,
650
+ ):
651
+ super().__init__()
652
+
653
+ self.attn = MultiheadAttention(
654
+ n_ctx,
655
+ width,
656
+ heads,
657
+ )
658
+ self.ln_1 = LayerNorm(width)
659
+ self.mlp = MLP(width)
660
+ self.ln_2 = LayerNorm(width)
661
+
662
+ def forward(self, x: torch.Tensor):
663
+ x = x + self.attn(self.ln_1(x))
664
+ x = x + self.mlp(self.ln_2(x))
665
+ return x
666
+
667
+
668
+
669
+ class LayerNorm(nn.LayerNorm):
670
+ """
671
+ Implementation that supports fp16 inputs but fp32 gains/biases.
672
+ """
673
+
674
+ def forward(self, x: torch.Tensor):
675
+ return super().forward(x.float()).to(x.dtype)
676
+
677
+
678
+ class Transformer(nn.Module):
679
+ def __init__(
680
+ self,
681
+ n_ctx: int,
682
+ width: int,
683
+ layers: int,
684
+ heads: int,
685
+ ):
686
+ super().__init__()
687
+ self.n_ctx = n_ctx
688
+ self.width = width
689
+ self.layers = layers
690
+ self.resblocks = nn.ModuleList(
691
+ [
692
+ ResidualAttentionBlock(
693
+ n_ctx,
694
+ width,
695
+ heads,
696
+ )
697
+ for _ in range(layers)
698
+ ]
699
+ )
700
+
701
+ def forward(self, x: torch.Tensor):
702
+ for block in self.resblocks:
703
+ x = block(x)
704
+ return x
705
+
706
+
707
+ class FrozenCLIPImageEmbedder(AbstractEncoder):
708
+ """Uses the CLIP transformer encoder for text (from Hugging Face)"""
709
+
710
+ def __init__(self, version='openai/clip-vit-large-patch14'):
711
+ super().__init__()
712
+ self.transformer = CLIPVisionModel.from_pretrained(version)
713
+ self.final_ln = LayerNorm(1024)
714
+ self.mapper = Transformer(
715
+ 1,
716
+ 1024,
717
+ 5,
718
+ 1,
719
+ )
720
+ self.shallow_index = 12
721
+ self.freeze()
722
+
723
+ def freeze(self):
724
+ self.transformer = self.transformer.eval()
725
+ for param in self.parameters():
726
+ param.requires_grad = False
727
+ for param in self.mapper.parameters():
728
+ param.requires_grad = True
729
+ for param in self.final_ln.parameters():
730
+ param.requires_grad = True
731
+
732
+ def forward(self, image):
733
+ outputs = self.transformer(pixel_values=image,
734
+ output_hidden_states=True)
735
+ all_hidden = outputs.hidden_states
736
+ last_hidden = outputs.last_hidden_state # b,257,1024
737
+ global_emb = last_hidden
738
+ gz = self.transformer.vision_model.post_layernorm(global_emb)
739
+ gz = self.mapper(gz)
740
+ z = self.final_ln(gz) # b,257,1024
741
+ z_local = z[:, 1:, :]
742
+ z_globaal = z[:, 0:1, :]
743
+ # get the outputs from shallow layer
744
+ # lz = all_hidden[self.shallow_index][:,1:] # b,256,1024
745
+ # z = torch.cat([gz, lz], dim=1)
746
+ return z_local, z_globaal
747
+
748
+ def encode(self, image):
749
+ return self(image)