add: replicate image generation

.gitignore

@@ -1,7 +1,143 @@
+# .gitignore
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
 *$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#  According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#  However, in case of collaboration, if having platform-specific dependencies or dependencies
+#  having no cross-platform support, pipenv may install dependencies that don't work, or not
+#  install all needed dependencies.
+# Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyderworkspace
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
 .idea/
 
+# Local logs
 *.log

app.py

@@ -1,3 +1,7 @@
+import cv2
+import numpy as np
+import requests
+from PIL.ImageOps import grayscale
 from fastapi.encoders import jsonable_encoder
 from src.utils import supabaseGetPublicURL, deductAndTrackCredit, returnBytesData
 from fastapi import FastAPI, File, UploadFile, Header, HTTPException, Form, Depends
@@ -12,13 +16,10 @@ import pandas as pd
 import base64
 import os
 from pydantic import BaseModel
-from gradio_client import Client, file
-from huggingface_hub import login
+import replicate
 
-HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
-login(HUGGINGFACE_TOKEN)
+from txt import jewellery_mask
 
-client = Client("techconsptr/CTO_TCP_V1")
 pipeline = Pipeline()
 app = FastAPI(title="Magical Mirror Web Plugin")
 
@@ -29,8 +30,8 @@ app.add_middleware(
     allow_methods=["*"],
     allow_headers=["*"],
 )
-url: str = os.environ["SUPABASE_URL"]
-key: str = os.environ["SUPABASE_KEY"]
+url: str = os.getenv("SUPABASE_URL")
+key: str = os.getenv("SUPABASE_KEY")
 supabase: Client = create_client(url, key)
 
 
@@ -40,14 +41,68 @@ class ClothingTryOnEntity(BaseModel):
 
 
 @app.post("/clothingTryOn")
-async def clothing_try_on(request: ClothingTryOnEntity):
-    result = client.predict(
-        input_image_base64=request.image,
-        mask_image_base64=request.mask,
-        api_name="/clothing_try_on_base64"
+async def clothing_try_on(image: UploadFile = File(...),
+                          mask: UploadFile = File(...)):
+    image_bytes = await image.read()
+    mask_bytes = await mask.read()
+    image, mask = Image.open(BytesIO(image_bytes)).convert("RGB"), Image.open(
+        BytesIO(mask_bytes)).convert("RGB")
+
+    actual_image = image.copy()
+
+    jewellery_mask = Image.fromarray(
+        np.bitwise_and(np.array(mask), np.array(image))
+    )
+
+    arr_orig = np.array(grayscale(mask))
+
+    image = cv2.inpaint(np.array(image), arr_orig, 15, cv2.INPAINT_TELEA)
+    image = Image.fromarray(image).resize((512, 512))
+
+    arr = arr_orig.copy()
+    mask_y = np.where(arr == arr[arr != 0][0])[0][0]
+    arr[mask_y:, :] = 255
+
+    mask = Image.fromarray(arr).resize((512, 512))
+
+    mask_img_base_64, act_img_base_64 = BytesIO(), BytesIO()
+    mask.save(mask_img_base_64, format="WEBP")
+    image.save(act_img_base_64, format="WEBP")
+    mask_bytes_ = base64.b64encode(mask_img_base_64.getvalue()).decode("utf-8")
+    image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
+
+    mask_data_uri = f"data:image/webp;base64,{mask_bytes_}"
+    image_data_uri = f"data:image/webp;base64,{image_bytes_}"
+
+    input = {
+        "mask": mask_data_uri,
+        "image": image_data_uri,
+        "prompt": "South Indian Saree, properly worn, natural setting, elegant, natural look, neckline without jewellery, simple",
+        "negative_prompt": "necklaces, jewellery, jewelry, necklace, neckpiece, garland, chain, neck wear, jewelled neck, jeweled neck, necklace on neck, jewellery on neck, accessories, watermark, text, changed background, wider body, narrower body, bad proportions, extra limbs, mutated hands, changed sizes, altered proportions, unnatural body proportions, blury, ugly",
+        "num_inference_steps": 25
+    }
+
+    output = replicate.run(
+        "stability-ai/stable-diffusion-inpainting:95b7223104132402a9ae91cc677285bc5eb997834bd2349fa486f53910fd68b3",
+        input=input
     )
+    image_url = output[0]
+
+    response = requests.get(image_url)
+    output_image = Image.open(BytesIO(response.content)).resize(actual_image.size)
+
+    result = Image.fromarray(np.bitwise_or(np.array(output_image), np.array(jewellery_mask)))
+
+    in_mem_file = BytesIO()
+    result.save(in_mem_file, format="WEBP", quality=85)
+    base_64_output = base64.b64encode(in_mem_file.getvalue()).decode('utf-8')
 
-    return JSONResponse(content=result, status_code=200)
+    response = {
+        "output": f"data:image/WEBP;base64,{base_64_output}",
+        'code': 200
+    }
+
+    return JSONResponse(content=response, status_code=200)
 
 
 @app.post("/productData/{storeId}")
@@ -248,3 +303,9 @@ async def necklace_try_on_id(necklace_try_on_id: NecklaceTryOnIDEntity = Depends
 
         else:
             return JSONResponse(content=response)
+
+
+if __name__ == '__main__':
+    import uvicorn
+
+    uvicorn.run(app, host="localhost", port=8000)

requirements.txt

@@ -77,6 +77,5 @@ uvicorn==0.30.1
 uvloop==0.19.0
 watchfiles==0.22.0
 websockets==12.0
-huggingface_hub
-gradio_client
+replicate
 -e .