Update app.py

app.py

@@ -229,11 +229,10 @@ h3{margin-top: 0}
 #component-1{text-align:center}
 .main_ui_logged_out{opacity: 0.3; pointer-events: none}
 .tabitem{border: 0px}
+#cost_preview_info{padding: .5em}
 """
 
-def swap_visibilty(profile: Union[gr.OAuthProfile, None], oauth_token: Union[gr.OAuthToken, None]):
-    user = whoami(oauth_token.token)
-    print(user)
+def swap_visibilty(profile: Union[gr.OAuthProfile, None]):
     if is_spaces:
         if profile is None:
             return gr.update(elem_classes=["main_ui_logged_out"])
@@ -242,14 +241,23 @@ def swap_visibilty(profile: Union[gr.OAuthProfile, None], oauth_token: Union[gr.
     else:
         return gr.update(elem_classes=["main_ui_logged_in"])
 
-def update_pricing(steps):
-    seconds_per_iteration = 7.54
-    total_seconds = (steps * seconds_per_iteration) + 240
-    cost_per_second = 0.80/60/60
-    cost = round(cost_per_second * total_seconds, 2)
-    cost_preview = f'''To train this LoRA, a paid L4 GPU will be hooked under the hood during training and then removed once finished.
-    ## Estimated to cost <b>< US$ {str(cost)}</b> for {round(int(total_seconds)/60, 2)} minutes with your current train settings <small>({int(steps)} iterations at {seconds_per_iteration}s/it)</small>'''
-    return gr.update(visible=True), cost_preview
+def update_pricing(steps, oauth_token: Union[gr.OAuthToken, None]):
+    if(oauth_token and is_spaces):
+        user = whoami(oauth_token.token)
+        seconds_per_iteration = 7.54
+        total_seconds = (steps * seconds_per_iteration) + 240
+        cost_per_second = 0.80/60/60
+        cost = round(cost_per_second * total_seconds, 2)
+        cost_preview = f'''To train this LoRA, a paid L4 GPU will be hooked under the hood during training and then removed once finished.
+        ### Estimated to cost <b>< US$ {str(cost)}</b> for {round(int(total_seconds)/60, 2)} minutes with your current train settings <small>({int(steps)} iterations at {seconds_per_iteration}s/it)</small>'''
+        if(user["canPay"]):
+            return gr.update(visible=True), cost_preview, gr.update(visible=False), gr.update(visible=True)
+        else:
+            pay_disclaimer = f'''<b>## ⚠️ {user.name}, your account doesn't have a payment method. Set one up <a href='https://huggingface.co/settings/billing/payment' target='_blank'>here</a> and come back here to train your LoRA<br>'''
+            return gr.update(visible=True), 
+        return gr.update(visible=True), pay_disclaimer+cost_preview, gr.update(visible=True), gr.update(visible=False)
+    else:
+        return gr.update(visible=False), "", gr.update(visible=False), gr.update(visible=True)
 
 with gr.Blocks(theme=theme, css=css) as demo:
     gr.Markdown(
@@ -330,13 +338,14 @@ with gr.Blocks(theme=theme, css=css) as demo:
                 sample_1 = gr.Textbox(label="Test prompt 1")
                 sample_2 = gr.Textbox(label="Test prompt 2")
                 sample_3 = gr.Textbox(label="Test prompt 3")
-            with gr.Group(visible=False) as cost_preview:
-                cost_preview_info = gr.Markdown()
+            with gr.Column(visible=False) as cost_preview:
+                cost_preview_info = gr.Markdown(elem_id="cost_preview_info")
+                payment_update = gr.Button("I have set up a payment method", visible=False)
             output_components.append(sample)
             output_components.append(sample_1)
             output_components.append(sample_2)
             output_components.append(sample_3)
-            start = gr.Button("Start training")
+            start = gr.Button("Start training", visible=False)
         progress_area = gr.Markdown("")
 
     with gr.Tab("Train on your device" if is_spaces else "Instructions"):
@@ -387,13 +396,13 @@ with gr.Blocks(theme=theme, css=css) as demo:
     ).then(
         update_pricing,
         inputs=[steps],
-        outputs=[cost_preview, cost_preview_info]
+        outputs=[cost_preview, cost_preview_info, payment_update, start]
     )
-
-    steps.change(
-        update_pricing,
+    gr.on(
+        triggers=[steps.change, payment_update.click],
+        fn=update_pricing,
         inputs=[steps],
-        outputs=[cost_preview, cost_preview_info]
+        outputs=[cost_preview, cost_preview_info, payment_update, start]
     )
     
     start.click(fn=create_dataset, inputs=[images] + caption_list, outputs=dataset_folder).then(