Dhruv Diddi commited on
Commit
956a40b
1 Parent(s): 6356404

feat: add lang support

Browse files
Files changed (2) hide show
  1. app.py +8 -3
  2. requirements.txt +2 -1
app.py CHANGED
@@ -7,7 +7,8 @@ from PIL import Image
7
  import re
8
  import os
9
 
10
-
 
11
 
12
  auth_token = os.getenv("auth_token")
13
  model_id = "CompVis/stable-diffusion-v1-4"
@@ -15,10 +16,14 @@ device = "cpu"
15
  pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=auth_token)
16
  pipe = pipe.to(device)
17
 
18
- def infer(prompt, samples, steps, scale, seed):
 
 
 
 
19
  generator = torch.Generator(device=device).manual_seed(seed)
20
  images_list = pipe(
21
- [prompt] * samples,
22
  num_inference_steps=steps,
23
  guidance_scale=scale,
24
  generator=generator,
 
7
  import re
8
  import os
9
 
10
+ from libretranslatepy import LibreTranslateAPI
11
+ lt = LibreTranslateAPI("https://translate.argosopentech.com/")
12
 
13
  auth_token = os.getenv("auth_token")
14
  model_id = "CompVis/stable-diffusion-v1-4"
 
16
  pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=auth_token)
17
  pipe = pipe.to(device)
18
 
19
+ def infer(prompt, samples, steps, scale, seed):
20
+ lang_detected = lt.detect(prompt)["language]"
21
+ print(lang_detected)
22
+ translated = lt.translate(prompt, lang_detected, "en")
23
+ print(translated)
24
  generator = torch.Generator(device=device).manual_seed(seed)
25
  images_list = pipe(
26
+ [translated] * samples,
27
  num_inference_steps=steps,
28
  guidance_scale=scale,
29
  generator=generator,
requirements.txt CHANGED
@@ -10,4 +10,5 @@ scipy
10
  scikit-image
11
  diffusers==0.3.0
12
  transformers
13
- ftfy
 
 
10
  scikit-image
11
  diffusers==0.3.0
12
  transformers
13
+ ftfy
14
+ libretranslatepy