Update app.py

app.py

@@ -14,12 +14,68 @@ from matplotlib import pyplot as plt
 from torchvision import transforms
 # from diffusers import DiffusionPipeline
 
+import io
+import multiprocessing
+import random
+import time
+import imghdr
+from pathlib import Path
+from typing import Union
+# from loguru import logger
+
+from lama_cleaner.model_manager import ModelManager
+from lama_cleaner.schema import Config
+
+try:
+    torch._C._jit_override_can_fuse_on_cpu(False)
+    torch._C._jit_override_can_fuse_on_gpu(False)
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    torch._C._jit_set_nvfuser_enabled(False)
+except:
+    pass
+    
+from lama_cleaner.helper import (
+    load_img,
+    numpy_to_bytes,
+    resize_max_size,
+)
+
+NUM_THREADS = str(multiprocessing.cpu_count())
+
+# fix libomp problem on windows https://github.com/Sanster/lama-cleaner/issues/56
+os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
+
+os.environ["OMP_NUM_THREADS"] = NUM_THREADS
+os.environ["OPENBLAS_NUM_THREADS"] = NUM_THREADS
+os.environ["MKL_NUM_THREADS"] = NUM_THREADS
+os.environ["VECLIB_MAXIMUM_THREADS"] = NUM_THREADS
+os.environ["NUMEXPR_NUM_THREADS"] = NUM_THREADS
+if os.environ.get("CACHE_DIR"):
+    os.environ["TORCH_HOME"] = os.environ["CACHE_DIR"]
+
+BUILD_DIR = os.environ.get("LAMA_CLEANER_BUILD_DIR", "app/build")
+
+
 from share_btn import community_icon_html, loading_icon_html, share_js
 
 HF_TOKEN_SD = os.environ.get('HF_TOKEN_SD')  or True
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+def diffuser_callback(i, t, latents):
+    pass
+
+model = ModelManager(
+        name='lama',
+        device=device,
+        hf_access_token=HF_TOKEN_SD,
+        sd_disable_nsfw=False,
+        sd_cpu_textencoder=True,
+        sd_run_local=True,
+        callback=diffuser_callback,
+    )
+    
+
 '''
 pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", dtype=torch.float16, revision="fp16", use_auth_token=auth_token).to(device)