thliang01 commited on
Commit
765b5c0
1 Parent(s): c27f33e

feat: add spaces import for zero gpu

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import torch
 
2
  from diffusers import StableDiffusionPipeline
3
  import gradio as gr
4
 
@@ -6,6 +7,9 @@ repo = "IDKiro/sdxs-512-0.9"
6
  seed = 42
7
  weight_type = torch.float16
8
 
 
 
 
9
  # Load model.
10
  pipe = StableDiffusionPipeline.from_pretrained(repo, torch_dtype=weight_type)
11
 
@@ -15,7 +19,7 @@ generator = pipe
15
  if torch.cuda.is_available():
16
  generator = generator.to("cuda")
17
 
18
-
19
  def generate(prompts):
20
  images = generator(list(prompts)).images
21
  return [images]
 
1
  import torch
2
+ import spaces
3
  from diffusers import StableDiffusionPipeline
4
  import gradio as gr
5
 
 
7
  seed = 42
8
  weight_type = torch.float16
9
 
10
+ zero = torch.Tensor([0]).cuda()
11
+ print(zero.device) # <-- 'cpu' 🤔
12
+
13
  # Load model.
14
  pipe = StableDiffusionPipeline.from_pretrained(repo, torch_dtype=weight_type)
15
 
 
19
  if torch.cuda.is_available():
20
  generator = generator.to("cuda")
21
 
22
+ @spaces.GPU(duration=120)
23
  def generate(prompts):
24
  images = generator(list(prompts)).images
25
  return [images]