Muhammadreza commited on
Commit
340f8fa
1 Parent(s): d88077a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -3,6 +3,7 @@ import torch
3
  import modin.pandas as pd
4
  import numpy as np
5
  from diffusers import DiffusionPipeline
 
6
 
7
  device = "cuda" if torch.cuda.is_available() else "cpu"
8
 
@@ -16,7 +17,8 @@ if torch.cuda.is_available():
16
  else:
17
  pipe = DiffusionPipeline.from_pretrained("mann-e/Mann-E_Dreams", use_safetensors=True)
18
  pipe = pipe.to(device)
19
-
 
20
  def genie (prompt, negative_prompt, steps, seed):
21
  generator = np.random.seed(0) if seed == 0 else torch.manual_seed(seed)
22
  int_image = pipe(prompt=prompt, negative_prompt=negative_prompt, generator=generator, num_inference_steps=steps, guidance_scale=0.0).images[0]
 
3
  import modin.pandas as pd
4
  import numpy as np
5
  from diffusers import DiffusionPipeline
6
+ import spaces
7
 
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
9
 
 
17
  else:
18
  pipe = DiffusionPipeline.from_pretrained("mann-e/Mann-E_Dreams", use_safetensors=True)
19
  pipe = pipe.to(device)
20
+
21
+ @spaces.GPU
22
  def genie (prompt, negative_prompt, steps, seed):
23
  generator = np.random.seed(0) if seed == 0 else torch.manual_seed(seed)
24
  int_image = pipe(prompt=prompt, negative_prompt=negative_prompt, generator=generator, num_inference_steps=steps, guidance_scale=0.0).images[0]