add multi lang

main.py

@@ -10,6 +10,11 @@ from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from pydantic import Field
 from diffusers import StableDiffusionPipeline
+from transformers import (
+    pipeline,
+    MBart50TokenizerFast,
+    MBartForConditionalGeneration,
+)
 
 app = FastAPI()
 
@@ -28,7 +33,9 @@ def index():
 @app.post("/draw", response_model=ItemOut)
 def draw(data: Data):
 	if data.member_secret != "" and data.member_secret == os.environ.get("MEMBER_SECRET"):
-		print(f"Is CUDA available: {torch.cuda.is_available()}")
+
+		device = "cuda" if torch.cuda.is_available() else "cpu"
+		device_dict = {"cuda": 0, "cpu": -1}
 
 		seedno = 0
 		if '_seed' in data.string:
@@ -50,13 +57,32 @@ def draw(data: Data):
 		# https://huggingface.co/docs/hub/spaces-sdks-docker-first-demo
 		# how to validation: https://qiita.com/bee2/items/75d9c0d7ba20e7a4a0e9
 		# https://github.com/huggingface/diffusers
+		# https://github.com/huggingface/diffusers/pull/1142
+
+		# Add language detection pipeline
+		language_detection_model_ckpt = "papluca/xlm-roberta-base-language-detection"
+		language_detection_pipeline = pipeline("text-classification",
+		                                       model=language_detection_model_ckpt,
+		                                       device=device_dict[device])
+
+		# 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)
 
 		model_id = 'stabilityai/stable-diffusion-2'
 
 		#pipe = StableDiffusionPipeline.from_pretrained(model_id)
-		pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)
+		pipe = StableDiffusionPipeline.from_pretrained(
+			model_id, 
+			custom_pipeline="multilingual_stable_diffusion",
+			detection_pipeline=language_detection_pipeline,
+			translation_model=trans_model,
+			translation_tokenizer=trans_tokenizer,
+			revision='fp16', 
+			torch_dtype=torch.float16
+		)
 		pipe.enable_attention_slicing()	# reduce gpu usage
-		pipe = pipe.to('cuda')
+		pipe = pipe.to(device)
 
 		generator = torch.Generator("cuda").manual_seed(seed)
 		images = pipe(prompt, negative_prompt=n_prompt, guidance_scale=7.5, generator=generator, num_images_per_prompt=1).images