pengdaqian commited on
Commit
2a589fd
1 Parent(s): 2dfa0cd
Files changed (2) hide show
  1. app.py +20 -5
  2. requirements.txt +6 -1
app.py CHANGED
@@ -11,6 +11,16 @@ from share_btn import community_icon_html, loading_icon_html, share_js
11
  word_list_dataset = load_dataset("Gustavosta/Stable-Diffusion-Prompts")
12
  word_list = word_list_dataset["train"]['Prompt']
13
 
 
 
 
 
 
 
 
 
 
 
14
  is_gpu_busy = False
15
  def infer(prompt, negative, scale):
16
  global is_gpu_busy
@@ -19,12 +29,17 @@ def infer(prompt, negative, scale):
19
  raise gr.Error("Unsafe content found. Please try again with different prompts.")
20
 
21
  images = []
22
- url = os.getenv('JAX_BACKEND_URL')
23
  payload = {'prompt': prompt, 'negative_prompt': negative, 'guidance_scale': scale}
24
- images_request = requests.post(url, json = payload)
25
- for image in images_request.json()["images"]:
26
- image_b64 = (f"data:image/jpeg;base64,{image}")
27
- images.append(image_b64)
 
 
 
 
 
28
 
29
  return images
30
 
 
11
  word_list_dataset = load_dataset("Gustavosta/Stable-Diffusion-Prompts")
12
  word_list = word_list_dataset["train"]['Prompt']
13
 
14
+
15
+ from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
16
+ import torch
17
+
18
+ model_id = "stabilityai/stable-diffusion-2-1-base"
19
+
20
+ scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
21
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16)
22
+
23
+
24
  is_gpu_busy = False
25
  def infer(prompt, negative, scale):
26
  global is_gpu_busy
 
29
  raise gr.Error("Unsafe content found. Please try again with different prompts.")
30
 
31
  images = []
32
+ generator = torch.manual_seed(68781279)
33
  payload = {'prompt': prompt, 'negative_prompt': negative, 'guidance_scale': scale}
34
+ image = pipe(prompt=prompt,
35
+ negative_prompt=negative,
36
+ guidance_scale=scale,
37
+ num_inference_steps=15,
38
+ generator=generator,
39
+ height=512,
40
+ width=768).images[0]
41
+
42
+ images.append(image)
43
 
44
  return images
45
 
requirements.txt CHANGED
@@ -1 +1,6 @@
1
- python-dotenv
 
 
 
 
 
 
1
+ python-dotenv
2
+ diffusers
3
+ transformers
4
+ accelerate
5
+ scipy
6
+ safetensors