Spaces:
Runtime error
Runtime error
Dont allow esd-u if less than 40gb of gpu memory
Browse files
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=
|
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'
|