change model id

main.py

@@ -9,12 +9,11 @@ from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from pydantic import Field
-from diffusers import StableDiffusionPipeline
+from diffusers import DiffusionPipeline
 from transformers import (
     pipeline,
     MBart50TokenizerFast,
     MBartForConditionalGeneration,
-    CLIPFeatureExtractor
 )
 
 app = FastAPI()
@@ -69,19 +68,16 @@ 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'
+		model_id = 'CompVis/stable-diffusion-v1-4'
 
 		#pipe = StableDiffusionPipeline.from_pretrained(model_id)
-		pipe = StableDiffusionPipeline.from_pretrained(
+		pipe = DiffusionPipeline.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,
-			safety_checker=null_safety,
 			revision='fp16', 
 			torch_dtype=torch.float16
 		)
@@ -119,7 +115,3 @@ def image_grid(imgs, rows, cols):
     for i, img in enumerate(imgs):
         grid.paste(img, box=(i%cols*w, i//cols*h))
     return grid
-
-def null_safety(images, **kwargs):
-    return images, False
-