update docs

app.py

@@ -6,6 +6,9 @@ import json
 import os
 import random
 import tempfile
+import shortuuid
+from apscheduler.schedulers.background import BackgroundScheduler
+import shutil
 
 import gradio as gr
 import torch
@@ -105,15 +108,18 @@ def upload_stage1_result(stage1_param_path: str, stage1_result_path: str,
     if not UPLOAD_REPO_ID:
         return
     try:
-        random_folder = random.randint(0,1000)
-        hf_api.upload_file(path_or_fileobj=stage1_param_path,
-                           path_in_repo=f'stage1_params/{random_folder}/{save_name}.json',
-                           repo_id=UPLOAD_REPO_ID,
-                           repo_type='dataset')
-        hf_api.upload_file(path_or_fileobj=stage1_result_path,
-                           path_in_repo=f'stage1_results/{random_folder}/{save_name}.pth',
-                           repo_id=UPLOAD_REPO_ID,
-                           repo_type='dataset')
+        folder_params = "tmp/results/stage1_params"
+        folder_results = "tmp/results/stage1_results"
+        
+        path_params = f"{folder_params}/{save_name}.json"
+        path_results = f"{folder_results}/{save_name}.pth"
+        
+        os.makedirs(folder_params, exist_ok=True)
+        os.makedirs(folder_results, exist_ok=True)
+        
+        shutil.copy(stage1_param_path, path_params)
+        shutil.copy(stage1_result_path, path_results)
+
     except Exception as e:
         print(e)
 
@@ -140,18 +146,20 @@ def upload_stage2_info(stage1_param_file_hash_name: str,
         param_file.write(json.dumps(stage2_params))
     stage2_param_file_hash_name = get_param_file_hash_name(param_file.name)
     save_name = f'{stage1_param_file_hash_name}_{stage2_param_file_hash_name}'
-
+    
     try:
-        random_folder = random.randint(0,1000)
-        hf_api.upload_file(path_or_fileobj=param_file.name,
-                           path_in_repo=f'stage2_params/{random_folder}/{save_name}.json',
-                           repo_id=UPLOAD_REPO_ID,
-                           repo_type='dataset')
+        folder_params = "tmp/results/stage2_params"
+
+        os.makedirs(folder_params, exist_ok=True)
+        path_params = f"{folder_params}/{save_name}.json"
+        shutil.copy(param_file.name, path_params)
+        
         if UPLOAD_RESULT_IMAGE:
-            hf_api.upload_file(path_or_fileobj=stage2_output_path,
-                               path_in_repo=f'stage2_results/{random_folder}/{save_name}.png',
-                               repo_id=UPLOAD_REPO_ID,
-                               repo_type='dataset')
+            folder_results = "tmp/results/stage2_results"
+            os.makedirs(folder_results, exist_ok=True)
+            path_results = f"{folder_results}/{save_name}.png"
+            shutil.copy(stage2_output_path, path_results)
+
     except Exception as e:
         print(e)
 
@@ -188,17 +196,16 @@ def upload_stage2_3_info(stage1_param_file_hash_name: str,
     save_name = f'{stage1_param_file_hash_name}_{stage2_3_param_file_hash_name}'
 
     try:
-        random_folder = random.randint(0,1000)
-        hf_api.upload_file(path_or_fileobj=param_file.name,
-                           path_in_repo=f'stage2_3_params/{random_folder}/{save_name}.json',
-                           repo_id=UPLOAD_REPO_ID,
-                           repo_type='dataset')
+        folder_params = "tmp/results/stage2_3_params"
+        os.makedirs(folder_params, exist_ok=True)
+        path_params = f"{folder_params}/{save_name}.json"
+        shutil.copy(param_file.name, path_params)
+
         if UPLOAD_RESULT_IMAGE:
-            hf_api.upload_file(
-                path_or_fileobj=stage2_3_output_path,
-                path_in_repo=f'stage2_3_results/{random_folder}/{save_name}.png',
-                repo_id=UPLOAD_REPO_ID,
-                repo_type='dataset')
+            folder_results = "tmp/results/stage2_3_results"
+            os.makedirs(folder_results, exist_ok=True)
+            path_results = f"{folder_results}/{save_name}.png"
+            shutil.copy(stage2_3_output_path, path_results)
     except Exception as e:
         print(e)
 
@@ -221,6 +228,22 @@ def show_gallery_view() -> tuple[dict, dict]:
 def show_upscaled_view() -> tuple[dict, dict]:
     return _update_result_view(False)
 
+def upload_files():
+    """Zips files and uploads to dataset. Local data is deleted
+    """
+    if os.path.exists("tmp/results") and os.path.isdir("tmp/results"):
+        try:
+            random_folder = random.randint(0,1000)
+            shutil.make_archive("tmp/results", 'zip', "tmp/results")
+            hf_api.upload_file(
+                path_or_fileobj="tmp/results.zip",
+                path_in_repo=f"{random_folder}/results_{shortuuid.uuid()}.zip",
+                repo_id=UPLOAD_REPO_ID,
+                repo_type="dataset",
+            )
+            shutil.rmtree("tmp/results")
+        except Exception as e:
+            print(e)
 
 examples = [
     'high quality dslr photo, a photo product of a lemon inspired by natural and organic materials, wooden accents, intricately decorated with glowing vines of led lights, inspired by baroque luxury',
@@ -239,6 +262,8 @@ examples = [
 ]
 
 with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(TITLE)
+    gr.Markdown(DESCRIPTION)
     with gr.Box():
         with gr.Row(elem_id='prompt-container').style(equal_height=True):
             with gr.Column():
@@ -249,6 +274,13 @@ with gr.Blocks(css='style.css') as demo:
                     placeholder='Enter your prompt',
                     elem_id='prompt-text-input',
                 ).style(container=False)
+                negative_prompt = gr.Text(
+                    label='Negative prompt',
+                    show_label=False,
+                    max_lines=1,
+                    placeholder='Enter a negative prompt',
+                    elem_id='negative-prompt-text-input',
+                ).style(container=False)
             generate_button = gr.Button('Generate').style(full_width=False)
 
         with gr.Column() as gallery_view:
@@ -274,7 +306,12 @@ with gr.Blocks(css='style.css') as demo:
                               interactive=False,
                               elem_id='upscaled-image').style(height=640)
             back_to_selection_button = gr.Button('Back to selection')
-
+            with gr.Group(elem_id="share-btn-container"):
+                community_icon = gr.HTML(community_icon_html)
+                loading_icon = gr.HTML(loading_icon_html)
+                share_button = gr.Button(
+                    "Share to community", elem_id="share-btn")
+                share_button.click(None, [], [], _js=share_js)
         with gr.Accordion('Advanced options',
                           open=False,
                           visible=SHOW_ADVANCED_OPTIONS):
@@ -385,12 +422,12 @@ with gr.Blocks(css='style.css') as demo:
     gr.HTML(FOOTER)
     stage1_inputs = [
         prompt,
-        "",
-        7,
-        2,
-        10,
-        "smart50",
-        25,
+        negative_prompt,
+        seed_1,
+        num_images,
+        guidance_scale_1,
+        custom_timesteps_1,
+        num_inference_steps_1,
     ]
     stage1_outputs = [
         gallery,
@@ -441,6 +478,49 @@ with gr.Blocks(css='style.css') as demo:
         queue=False,
     )
 
+    negative_prompt.submit(
+        fn=randomize_seed_fn,
+        inputs=[seed_1, randomize_seed_1],
+        outputs=seed_1,
+        queue=False,
+    ).then(
+        fn=lambda: -1,
+        outputs=selected_index_for_stage2,
+        queue=False,
+    ).then(
+        fn=show_gallery_view,
+        outputs=[
+            gallery_view,
+            upscale_view,
+        ],
+        queue=False,
+    ).then(
+        fn=update_upscale_button,
+        inputs=selected_index_for_stage2,
+        outputs=[
+            upscale_button,
+            upscale_to_256_button,
+        ],
+        queue=False,
+    ).then(
+        fn=model.run_stage1,
+        inputs=stage1_inputs,
+        outputs=stage1_outputs,
+    ).success(
+        fn=get_param_file_hash_name,
+        inputs=stage1_param_path,
+        outputs=stage1_param_file_hash_name,
+        queue=False,
+    ).then(
+        fn=upload_stage1_result,
+        inputs=[
+            stage1_param_path,
+            stage1_result_path,
+            stage1_param_file_hash_name,
+        ],
+        queue=False,
+    )
+
     generate_button.click(
         fn=randomize_seed_fn,
         inputs=[seed_1, randomize_seed_1],
@@ -553,7 +633,7 @@ with gr.Blocks(css='style.css') as demo:
         custom_timesteps_2,
         num_inference_steps_2,
         prompt,
-        "",
+        negative_prompt,
         seed_3,
         guidance_scale_3,
         num_inference_steps_3,
@@ -596,7 +676,7 @@ with gr.Blocks(css='style.css') as demo:
             custom_timesteps_2,
             num_inference_steps_2,
             prompt,
-            "",
+            negative_prompt,
             seed_3,
             guidance_scale_3,
             num_inference_steps_3,
@@ -612,5 +692,10 @@ with gr.Blocks(css='style.css') as demo:
         ],
         queue=False,
     )
+    
+    if UPLOAD_REPO_ID:
+        scheduler = BackgroundScheduler()
+        scheduler.add_job(func=upload_files, trigger="interval", seconds=60*20)
+        scheduler.start()
 
-demo.queue(api_open=True, max_size=MAX_QUEUE_SIZE).launch()
+demo.queue(api_open=False, max_size=MAX_QUEUE_SIZE).launch(debug=DEBUG)