add feature_extractor

Dockerfile

@@ -37,8 +37,7 @@ RUN python3 -m pip install --no-cache-dir --upgrade pip && \
         numpy \
         scipy \
         tensorboard \
-        transformers \
-        streamlit
+        transformers
 
 # Set the working directory to /code
 WORKDIR /code

main.py

@@ -4,17 +4,17 @@ import sys
 import re
 import random
 import torch
-import streamlit as st
 from PIL import Image
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from pydantic import Field
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionPipeline
 from transformers import (
     pipeline,
     MBart50TokenizerFast,
     MBartForConditionalGeneration,
+    CLIPFeatureExtractor
 )
 
 app = FastAPI()
@@ -69,16 +69,18 @@ def draw(data: Data):
 		# Add model for language translation
 		trans_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-one-mmt")
 		trans_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-one-mmt").to(device)
+		feature_extractor = CLIPFeatureExtractor.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K")
 
 		model_id = 'stabilityai/stable-diffusion-2'
 
 		#pipe = StableDiffusionPipeline.from_pretrained(model_id)
-		pipe = DiffusionPipeline.from_pretrained(
+		pipe = StableDiffusionPipeline.from_pretrained(
 			model_id, 
 			custom_pipeline="multilingual_stable_diffusion",
 			detection_pipeline=language_detection_pipeline,
 			translation_model=trans_model,
 			translation_tokenizer=trans_tokenizer,
+			feature_extractor=feature_extractor,
 			revision='fp16', 
 			torch_dtype=torch.float16
 		)