add translator

main.py

@@ -9,12 +9,8 @@ from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from pydantic import Field
-from diffusers import DiffusionPipeline
-from transformers import (
-    pipeline,
-    MBart50TokenizerFast,
-    MBartForConditionalGeneration,
-)
+from diffusers import StableDiffusionPipeline
+from googletrans import Translator
 
 app = FastAPI()
 
@@ -50,6 +46,10 @@ def draw(data: Data):
 		text = re.sub('^#', '', data.string)
 		text = re.sub('_seed[1-4]?', '', text)
 
+		translator = Translator()
+		translation = translator.translate(text, dest="en")
+		text = translation.text
+
 		# prompt = '(('+text+')) (( photograph )), highly detailed, sharp focus, 8k, 4k, (( photorealism )), detailed, saturated, portrait, 50mm, F/2.8, 1m away, ( global illumination, studio light, volumetric light ), ((( multicolor lights )))'
 		prompt = '(('+text+')) (( photograph )), highly detailed, sharp focus, 8k, 4k, (( photorealism )), detailed, saturated, portrait, 50mm, F/2.8, 1m away, ((( multicolor lights )))'
 		n_prompt = 'text, blurry, art, painting, rendering, drawing, sketch, (( ugly )), (( duplicate )), ( morbid ), (( mutilated )), ( mutated ), ( deformed ), ( disfigured ), ( extra limbs ), ( malformed limbs ), ( missing arms ), ( missing legs ), ( extra arms ), ( extra legs ), ( fused fingers ), ( too many fingers ), long neck, low quality, worst quality'
@@ -57,34 +57,15 @@ 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 = 'CompVis/stable-diffusion-v1-4'
+		model_id = 'stabilityai/stable-diffusion-2'
 
 		#pipe = StableDiffusionPipeline.from_pretrained(model_id)
-		pipe = DiffusionPipeline.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 = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)
 		pipe.enable_attention_slicing()	# reduce gpu usage
 		pipe = pipe.to(device)
 
-		generator = torch.Generator("cuda").manual_seed(seed)
+		generator = torch.Generator(device).manual_seed(seed)
 		images = pipe(prompt, negative_prompt=n_prompt, guidance_scale=7.5, generator=generator, num_images_per_prompt=1).images
 
 		if 0 < int(seedno) < 5:

requirements.txt

@@ -2,4 +2,5 @@ fastapi==0.74.*
 requests==2.27.*
 sentencepiece==0.1.*
 uvicorn[standard]==0.17.*
-diffusers[torch]==0.13.*
+diffusers[torch]==0.13.*
+googletrans==4.0.0-rc1