rnlduatm commited on
Commit
498d3e0
ยท
1 Parent(s): fc51b17

Update space

Browse files
Files changed (1) hide show
  1. model/waifu.py +8 -11
model/waifu.py CHANGED
@@ -23,17 +23,14 @@
23
  from diffusers import StableDiffusionPipeline
24
  import torch
25
 
26
-
27
- def generate_waifu(prompt: str, model_id):
28
- model_id = "hakurei/waifu-diffusion"
29
- pipe = StableDiffusionPipeline.from_pretrained(
30
- model_id,
31
- torch_dtype=torch.float32, # CPU๋Š” float32๋งŒ ์ง€์›
32
- use_safetensors=False
33
- # revision=None # revision="fp16" ์‚ญ์ œ
34
- ).to("cpu") # ๋ฐ˜๋“œ์‹œ CPU๋กœ ์ง€์ •
35
-
36
-
37
  image = pipe(prompt=prompt).images[0]
38
  return image
39
 
 
23
  from diffusers import StableDiffusionPipeline
24
  import torch
25
 
26
+ model_id = "hakurei/waifu-diffusion"
27
+ pipe = StableDiffusionPipeline.from_pretrained(
28
+ model_id,
29
+ torch_dtype=torch.float32,
30
+ use_safetensors=False
31
+ ).to("cpu") # CPU๋กœ ๋ช…์‹œ
32
+
33
+ def generate_waifu(prompt: str):
 
 
 
34
  image = pipe(prompt=prompt).images[0]
35
  return image
36