adapt handler from https://huggingface.co/spaces/cfahlgren1/Emoji-Generator-by-fofr

handler.py

@@ -1,23 +1,37 @@
-from typing import  Dict, List, Any
+from typing import Dict, List, Any
 import torch
 from torch import autocast
-from diffusers import StableDiffusionPipeline
+from huggingface_hub import hf_hub_download
+from diffusers import DiffusionPipeline
 import base64
 from io import BytesIO
+from cog_sdxl.dataset_and_utils import TokenEmbeddingsHandler
 
 
-# set device
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("device ~>", device)
 
-if device.type != 'cuda':
-    raise ValueError("need to run on GPU")
 
-class EndpointHandler():
+class EndpointHandler:
     def __init__(self, path=""):
-        # load the optimized model
-        self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
-        self.pipe = self.pipe.to(device)
+        print("path ~>", path)
 
+        self.pipe = DiffusionPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            torch_dtype=torch.float16 if device.type == "cuda" else None,
+            variant="fp16",
+        ).to(device)
+
+        self.pipe.load_lora_weights("SvenN/sdxl-emoji", weight_name="lora.safetensors")
+
+        text_encoders = [self.pipe.text_encoder, self.pipe.text_encoder_2]
+        tokenizers = [self.pipe.tokenizer, self.pipe.tokenizer_2]
+
+        embedding_path = hf_hub_download(
+            repo_id="SvenN/sdxl-emoji", filename="embeddings.pti", repo_type="model"
+        )
+        embhandler = TokenEmbeddingsHandler(text_encoders, tokenizers)
+        embhandler.load_embeddings(embedding_path)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
@@ -28,15 +42,25 @@ class EndpointHandler():
             A :obj:`dict`:. base64 encoded image
         """
         inputs = data.pop("inputs", data)
-        
-        # run inference pipeline
-        with autocast(device.type):
-            image = self.pipe(inputs, guidance_scale=7.5)["sample"][0]  
-            
+
+        # Automatically add trigger tokens to the beginning of the prompt
+        full_prompt = f"A <s0><s1> {inputs}"
+        images = self.pipe(
+            full_prompt,
+            cross_attention_kwargs={"scale": 0.8},
+        ).images
+        image = images[0]
         # encode image as base 64
         buffered = BytesIO()
         image.save(buffered, format="JPEG")
         img_str = base64.b64encode(buffered.getvalue())
 
         # postprocess the prediction
-        return {"image": img_str.decode()}
+        return {"image": img_str.decode()}
+
+
+if __name__ == "__main__":
+    handler = EndpointHandler()
+    print(handler)
+    output = handler({"inputs": "emoji of a tiger face, white background"})
+    print(output)

requirements.txt

@@ -0,0 +1,88 @@
+accelerate==0.31.0
+aiofiles==23.2.1
+altair==5.3.0
+annotated-types==0.7.0
+anyio==4.4.0
+attrs==23.2.0
+certifi==2024.6.2
+charset-normalizer==3.3.2
+click==8.1.7
+contourpy==1.2.1
+cycler==0.12.1
+diffusers==0.29.1
+dnspython==2.6.1
+email_validator==2.2.0
+exceptiongroup==1.2.1
+fastapi==0.111.0
+fastapi-cli==0.0.4
+ffmpy==0.3.2
+filelock==3.15.4
+fonttools==4.53.0
+fsspec==2024.6.0
+gradio==4.36.1
+gradio_client==1.0.1
+h11==0.14.0
+httpcore==1.0.5
+httptools==0.6.1
+httpx==0.27.0
+huggingface-hub==0.23.4
+idna==3.7
+importlib_metadata==7.2.0
+importlib_resources==6.4.0
+Jinja2==3.1.4
+jsonschema==4.22.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.0
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.3
+numpy==1.23.5
+orjson==3.10.5
+packaging==24.1
+pandas==2.2.2
+peft==0.11.1
+pillow==10.3.0
+psutil==5.9.8
+pydantic==2.7.4
+pydantic_core==2.18.4
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.9
+pytz==2024.1
+PyYAML==6.0.1
+referencing==0.35.1
+regex==2024.5.15
+requests==2.32.3
+rich==13.7.1
+rpds-py==0.18.1
+ruff==0.4.10
+safetensors==0.4.3
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+spaces==0.28.3
+starlette==0.37.2
+sympy==1.12.1
+tokenizers==0.19.1
+tomlkit==0.12.0
+toolz==0.12.1
+torch==2.2.0
+tqdm==4.66.4
+transformers==4.41.2
+typer==0.12.3
+typing_extensions==4.12.2
+tzdata==2024.1
+ujson==5.10.0
+urllib3==2.2.2
+uvicorn==0.30.1
+uvloop==0.19.0
+watchfiles==0.22.0
+websockets==11.0.3
+zipp==3.19.2