wjs0725 commited on
Commit
aca4c0c
1 Parent(s): 779331d

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -9,7 +9,7 @@ import argparse
9
  from einops import rearrange
10
  from fire import Fire
11
  from PIL import ExifTags, Image
12
-
13
 
14
  import torch
15
  import torch.nn.functional as F
@@ -75,6 +75,7 @@ class FluxEditor:
75
  self.ae.encoder.to(self.device)
76
 
77
  @torch.inference_mode()
 
78
  def edit(self, init_image, source_prompt, target_prompt, num_steps, inject_step, guidance, seed):
79
  torch.cuda.empty_cache()
80
  seed = None
 
9
  from einops import rearrange
10
  from fire import Fire
11
  from PIL import ExifTags, Image
12
+ import spaces
13
 
14
  import torch
15
  import torch.nn.functional as F
 
75
  self.ae.encoder.to(self.device)
76
 
77
  @torch.inference_mode()
78
+ @spaces.GPU(duration=30)
79
  def edit(self, init_image, source_prompt, target_prompt, num_steps, inject_step, guidance, seed):
80
  torch.cuda.empty_cache()
81
  seed = None