JadenFK commited on
Commit
98bb9c3
1 Parent(s): d5684d7

Dont allow esd-u if less than 40gb of gpu memory

Browse files
Files changed (1) hide show
  1. app.py +6 -1
app.py CHANGED
@@ -103,8 +103,13 @@ class Demo:
103
  label="Prompt to Erase",
104
  info="Prompt corresponding to concept to erase"
105
  )
 
 
 
 
 
106
  self.train_method_input = gr.Dropdown(
107
- choices=['ESD-x', 'ESD-u', 'ESD-self'],
108
  value='ESD-x',
109
  label='Train Method',
110
  info='Method of training'
 
103
  label="Prompt to Erase",
104
  info="Prompt corresponding to concept to erase"
105
  )
106
+
107
+ choices = ['ESD-x']
108
+ if torch.cuda.get_device_properties(0).total_memory * 1e-9 >= 40:
109
+ choices.append('ESD-u')
110
+
111
  self.train_method_input = gr.Dropdown(
112
+ choices=choices,
113
  value='ESD-x',
114
  label='Train Method',
115
  info='Method of training'