instantid

Author: apocas

.gitignore

@@ -164,6 +164,8 @@ projects/*
 !projects/.gitkeep
 models/*
 !models/.gitkeep
+checkpoints/*
+!checkpoints/.gitkeep
 uploads/*
 !uploads/.gitkeep
 .DS_Store

Makefile

@@ -1,6 +1,6 @@
 .PHONY: start
 start:
-	poetry run python main.py
+	poetry run uvicorn app.main:app --port 9000 --workers 4
 
 .PHONY: dev
 dev:

app/brain.py

@@ -20,6 +20,7 @@
 from app.llms.tools.dalle import DalleImage
 from app.llms.tools.describeimage import DescribeImage
 from app.llms.tools.drawimage import DrawImage
+from app.llms.tools.instantid import InstantID
 from app.llms.tools.refineimage import RefineImage
 from app.llms.tools.stablediffusion import StableDiffusionImage
 from app.model import Model
@@ -375,6 +376,7 @@ def entryVision(self, projectName, visionInput, isprivate, db: Session):
             RefineImage(),
             DrawImage(),
             DescribeImage(),
+            InstantID(),
         ]
 
         if isprivate:

app/llms/instantid/pipeline_stable_diffusion_xl_instantid.py

@@ -0,0 +1,87 @@
+import base64
+import io
+from diffusers.models import ControlNetModel
+from huggingface_hub import hf_hub_download
+
+import cv2
+import torch
+import numpy as np
+import random
+from PIL import Image
+from insightface.app import FaceAnalysis
+
+from app.llms.instantid.pipeline_stable_diffusion_xl_instantid import StableDiffusionXLInstantIDPipeline, draw_kps
+
+def instantid_worker(prompt, sharedmem):
+    try:
+        hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/config.json", local_dir="./checkpoints")
+        hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/diffusion_pytorch_model.safetensors", local_dir="./checkpoints")
+        hf_hub_download(repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir="./checkpoints")
+    except:
+        pass
+
+    img_data = base64.b64decode(sharedmem["input_image"])
+    face_image = Image.open(io.BytesIO(img_data))
+
+    prompt_default = ", (detailed) (intricate) (8k) (HDR) (cinematic lighting) (sharp focus)"
+    prompt = prompt + prompt_default
+
+    negative_prompt = "(lowres, low quality, worst quality:1.2), (text:1.2), watermark, (frame:1.2), deformed, ugly, deformed eyes, blur, out of focus, blurry, deformed cat, deformed, photo, anthropomorphic cat, monochrome, photo, pet collar, gun, weapon, blue, 3d, drones, drone, buildings in background, green"
+
+    DEFAULT_CUDA = "cuda"
+
+    app = FaceAnalysis(name='antelopev2', root='./', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+    app.prepare(ctx_id=0, det_size=(640, 640))
+
+    face_adapter = f'./checkpoints/ip-adapter.bin'
+    controlnet_path = f'./checkpoints/ControlNetModel'
+
+    controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
+
+    pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
+        "wangqixun/YamerMIX_v8", controlnet=controlnet, torch_dtype=torch.float16
+    )
+    pipe.to(DEFAULT_CUDA)
+
+    pipe.load_ip_adapter_instantid(face_adapter)
+
+    pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")
+    pipe.disable_lora()
+
+    face_image_cv2 = cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR)
+    height, width, _ = face_image_cv2.shape
+
+    face_info = app.get(face_image_cv2)
+    face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*x['bbox'][3]-x['bbox'][1])[-1]
+    face_emb = face_info['embedding']
+    face_kps = draw_kps(face_image, face_info['kps'])
+
+    control_mask = np.zeros([height, width, 3])
+    x1, y1, x2, y2 = face_info["bbox"]
+    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+    control_mask[y1:y2, x1:x2] = 255
+    control_mask = Image.fromarray(control_mask.astype(np.uint8))
+
+    pipe.set_ip_adapter_scale(0.8)
+
+    generator = torch.Generator(device=DEFAULT_CUDA).manual_seed(random.randint(0, np.iinfo(np.int32).max))
+
+    image = pipe(
+        prompt,
+        image_embeds=face_emb,
+        image=face_kps,
+        control_mask=control_mask,
+        num_inference_steps=50,
+        controlnet_conditioning_scale=0.8,
+        negative_prompt=negative_prompt,
+        generator=generator,
+        guide_scale=0,
+        height=height,
+        width=width,
+    ).images[0]
+
+    output_img_data = io.BytesIO()
+    image.save(output_img_data, format="JPEG")
+    image_base64 = base64.b64encode(output_img_data.getvalue()).decode('utf-8')
+
+    sharedmem["output_image"] = image_base64

app/llms/instantid/worker.py

@@ -0,0 +1,49 @@
+from torch.multiprocessing import Process, set_start_method, Manager
+
+from app.llms.instantid.worker import instantid_worker
+try:
+    set_start_method('spawn')
+except RuntimeError:
+    pass
+from langchain.tools import BaseTool
+from langchain.chains import LLMChain
+from langchain_community.chat_models import ChatOpenAI
+from langchain.prompts import PromptTemplate
+
+from typing import Optional
+from langchain.callbacks.manager import (
+    CallbackManagerForToolRun,
+)
+
+
+class InstantID(BaseTool):
+    name = "Avatar Generator"
+    description = "use this tool when you need to draw an avatar from an image and a descripton."
+    return_direct = True
+
+    def _run(self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
+        if run_manager.tags[0].boost == True:
+            llm = ChatOpenAI(temperature=0.9, model_name="gpt-3.5-turbo")
+            prompt = PromptTemplate(
+                input_variables=["image_desc"],
+                template="Generate a detailed prompt to generate an image based on the following description: {image_desc}",
+            )
+            chain = LLMChain(llm=llm, prompt=prompt)
+
+            fprompt = chain.run(query)
+        else:
+            fprompt = run_manager.tags[0].question
+
+        manager = Manager()
+        sharedmem = manager.dict()
+
+        sharedmem["input_image"] = run_manager.tags[0].image
+
+        p = Process(target=instantid_worker, args=(fprompt, sharedmem))
+        p.start()
+        p.join()
+
+        return {"type": "instantid", "image": sharedmem["output_image"], "prompt": fprompt}
+
+    async def _arun(self, query: str) -> str:
+        raise NotImplementedError("N/A")