use the fp16 revision of SD (#72)

Author: jeffra

mii/grpc_related/modelresponse_server.py

@@ -10,7 +10,6 @@
 import sys
 import time
 
-from torch import autocast
 from transformers import Conversation
 from mii.constants import GRPC_MAX_MSG_SIZE
 
@@ -75,8 +74,7 @@ def Txt2ImgReply(self, request, context):
         request = [r for r in request.request]
 
         start = time.time()
-        with autocast("cuda"):
-            response = self.inference_pipeline(request, **query_kwargs)
+        response = self.inference_pipeline(request, **query_kwargs)
         end = time.time()
 
         images_bytes = []

mii/models/providers/diffusers.py

@@ -1,10 +1,18 @@
 import os
+import torch
 
 
 def diffusers_provider(model_path, model_name, task_name, mii_config):
     from diffusers import DiffusionPipeline
     local_rank = int(os.getenv('LOCAL_RANK', '0'))
+
+    kwargs = {}
+    if mii_config.torch_dtype() == torch.half:
+        kwargs["torch_dtype"] = torch.float16
+        kwargs["revision"] = "fp16"
+
     pipeline = DiffusionPipeline.from_pretrained(model_name,
-                                                 use_auth_token=mii_config.hf_auth_token)
+                                                 use_auth_token=mii_config.hf_auth_token,
+                                                 **kwargs)
     pipeline = pipeline.to(f"cuda:{local_rank}")
     return pipeline