Optimize handler, update README

README.md

@@ -4,47 +4,30 @@ tags:
 - stable-diffusion
 - stable-diffusion-diffusers
 - endpoints-template
-duplicated_from: philschmid/stable-diffusion-v1-4-endpoints
+
+extra_gated_prompt: |-
+  This model is open access and available to all, with a CreativeML OpenRAIL-M license further specifying rights and usage.
+  The CreativeML OpenRAIL License specifies: 
+
+  1. You can't use the model to deliberately produce nor share illegal or harmful outputs or content 
+  2. CompVis claims no rights on the outputs you generate, you are free to use them and are accountable for their use which must not go against the provisions set in the license
+  3. You may re-distribute the weights and use the model commercially and/or as a service. If you do, please be aware you have to include the same use restrictions as the ones in the license and share a copy of the CreativeML OpenRAIL-M to all your users (please read the license entirely and carefully)
+  Please read the full license carefully here: https://huggingface.co/spaces/CompVis/stable-diffusion-license
+
+extra_gated_heading: Please read the LICENSE to access this model
 ---
-# Fork of [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4)
-> Stable Diffusion is a latent text-to-image diffusion model capable of generating photo-realistic images given any text input.
-> For more information about how Stable Diffusion functions, please have a look at [🤗's Stable Diffusion with 🧨Diffusers blog](https://huggingface.co/blog/stable_diffusion).
-For more information about the model, license and limitations check the original model card at [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4).
-### License (CreativeML OpenRAIL-M)
-The full license can be found here: https://huggingface.co/spaces/CompVis/stable-diffusion-license
----
-This repository implements a custom `handler` task for `text-to-image` for 🤗 Inference Endpoints. The code for the customized pipeline is in the [pipeline.py](https://huggingface.co/philschmid/stable-diffusion-v1-4-endpoints/blob/main/handler.py).
-There is also a [notebook](https://huggingface.co/philschmid/stable-diffusion-v1-4-endpoints/blob/main/create_handler.ipynb) included, on how to create the `handler.py`
-### expected Request payload
-```json
-{
-    "inputs": "A prompt used for image generation"
-}
-```
-below is an example on how to run a request using Python and `requests`.
-## Run Request 
-```python
-import json
-from typing import List
-import requests as r
-import base64
-from PIL import Image
-from io import BytesIO
-ENDPOINT_URL = ""
-HF_TOKEN = ""
-# helper decoder
-def decode_base64_image(image_string):
-  base64_image = base64.b64decode(image_string)
-  buffer = BytesIO(base64_image)
-  return  Image.open(buffer)
-def predict(prompt:str=None):
-    payload = {"inputs": code_snippet,"parameters": parameters}
-    response = r.post(
-        ENDPOINT_URL, headers={"Authorization": f"Bearer {HF_TOKEN}"}, json={"inputs": prompt}
-    )
-    resp = response.json()
-    return decode_base64_image(resp["image"])
-prediction = predict(
-    prompt="the first animal on the mars"
-)
-```
+# Stable Diffusion v1-5 Custom Inference
+This repo is for running diffusion custom inference endpoints with `prompts` and an optional `image` as inputs (Unlike normal text-to-image inference). To 
+achieve this goal, this repo implements a `handler.py` script. For more information regarding custom inference, please visit 
+this [link](https://huggingface.co/docs/inference-endpoints/guides/custom_handler).
+For more information about the model, license and limitations please check the original [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5)
+or diffusion [documentation](https://huggingface.co/docs/diffusers/index).
+
+### Local test custom handler
+To test custom inference locally, please run the following command:
+```commandline
+ python local_request.py --prompts="whale in the universe" --image="test_image.jpg"
+``` 
+**Note**: `--image` parameter is optional.
+
+

handler.py

@@ -1,7 +1,6 @@
 import base64
 from io import BytesIO
-from typing import Dict, List, Any
-import os
+from typing import Dict, Any
 
 import torch
 from PIL import Image
@@ -17,17 +16,15 @@ def decode_base64_image(image_string):
 
 class EndpointHandler:
     def __init__(self, path=""):
-        print(list(os.walk(".")))
-        self.pipe = StableDiffusionPipeline.from_pretrained("/repository/stable-diffusion-v1-5")
+        self.pipe = StableDiffusionPipeline.from_pretrained("/repository/stable-diffusion-v1-5",
+            torch_dtype=torch.float16, revision="fp16")
         self.pipe = self.pipe.to("cuda")
 
-    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+    def __call__(self, data: Any) -> Dict[str, str]:
         """
-        Args:
-            data (:obj:):
-                includes the input data and the parameters for the inference.
-        Return:
-            A :obj:`dict`:. base64 encoded image
+        Return predict value.
+        :param data: A dictionary contains `inputs` and optional `image` field.
+        :return: A dictionary with `image` field contains image in base64.
         """
         prompts = data.pop("inputs", None)
         encoded_image = data.pop("image", None)
@@ -35,11 +32,10 @@ class EndpointHandler:
         if encoded_image:
             init_image = decode_base64_image(encoded_image)
             init_image.thumbnail((768, 768))
-        image = self.pipe(prompts, init_image=init_image).images[0]
 
-        # encode image as base 64
+        image = self.pipe(prompts, init_image=init_image).images[0]
         buffered = BytesIO()
         image.save(buffered, format="png")
+        img_str = base64.b64encode(buffered.getvalue())
 
-        # post process the prediction
-        return {"image": buffered.getvalue()}
+        return {"image": img_str.decode()}