app.py

import gradio as gr
from gradio_modal import Modal
from civitai_api import (search_civitai, download_civitai, select_civitai_item, add_civitai_item, get_civitai_tag, select_civitai_all_item,
                           update_civitai_selection, update_civitai_checkbox, from_civitai_checkbox)
from civitai_constants import (TYPE, BASEMODEL, SORT, PERIOD, FILETYPE)
import time

css = """

.title { padding: 30px 00px; font-size: 3em; align-items: center; text-align: center; }

.info { align-items: center; text-align: center; }

.desc [src$='#float'] { float: right; margin: 20px; }

#modal-window{ position: absolute; left: 50%; top: 50%; transform: translate(-50%, -50%); max-height: inherit; }

.modal-container{ max-width:60vw; height:100vh; }

.prose.desc{ padding: 0px 20px; }

/* Video hover styles */

.media-container {

    position: relative;

    width: 100%;

    height: 100%;

}

.result{ padding: 16px; }



.media-container img,

.media-container video {

    width: 100%;

    height: 100%;

    object-fit: contain;

}



.media-container video {

    display: none;

    position: absolute;

    top: 0;

    left: 0;

    z-index: 2;

}



.gallery-item:hover .media-container video {

    display: block;

}



.gallery-item:hover .media-container img {

    opacity: 0.3;

}



.gallery-item {

    position: relative;

    cursor: pointer;

    overflow: hidden;

}

"""

# js = """
# (() => {
#     function setupGalleryHover() {
#         const gallery = document.getElementById('gallery');
#         if (!gallery) return;

#         function handleMediaContainer(container) {
#             const video = container.querySelector('video');
#             if (!video) return;
            
#             video.muted = true;
#             video.preload = "metadata";
            
#             container.addEventListener('mouseenter', () => {
#                 if (video.paused) {
#                     video.play().catch(() => {});
#                 }
#             });
            
#             container.addEventListener('mouseleave', () => {
#                 if (!video.paused) {
#                     video.pause();
#                     video.currentTime = 0;
#                 }
#             });
#         }

#         function setupContainers() {
#             const containers = gallery.querySelectorAll('.media-container');
#             containers.forEach(handleMediaContainer);
#         }

#         setupContainers();
        
#         const observer = new MutationObserver(setupContainers);
#         observer.observe(gallery, { childList: true, subtree: true });
#     }

#     if (document.readyState === 'loading') {
#         document.addEventListener('DOMContentLoaded', setupGalleryHover);
#     } else {
#         setupGalleryHover();
#     }
# })();
# """

common_sense = """

## The Common Sense Agreement



This tool was created in response to recent censorship and model deactivations on Civitai. This Space should be treated as a tool for artists and creators to easily backup and archive models -- this is NOT a tool to blindly scrape and download every asset your heart desires. For those that wish to do so, may do so on their own local machines. Users of the tool are responsible for their own actions - you have been warned.



By continuing on to the tool, you are assumed to have **acknowledged, understood, and are in agreeance** with the following common sense keypoints:



1. **Storage is not free.** You are responsible for understanding your storage limits and the costs that will be incurred if you exceed your tier's free allocation.

2. **Bandwidth is not free.** Misuse of the tool will likely be frowned upon by both Huggingface as well as Civitai, potentially leading to negative consequences for all. Don't be that guy. There are better ways if you must do so. 

3. **Not all models are created equally.** Respect the licensing agreements and the wishes of creators, both big and small. 

4. **Respect the FOSS community**, give back where you can, and be aware of your surroundings. Protect this industry as best you can, otherwise prepare for a wicked spanking when AGI is here.



### For those with common sense, please continue.

"""
header_md = """

# 🚨 Civitai Models & Workflow Conservation / Archiving Tool

### WIP: Expect changes and adjustments. Planned: S3 buckets, B2 storage, IPFS, magnet tools for torrents. 

---

* Provided by [AI Without Borders](https://huggingface.co/aiwithoutborders-xyz). Run locally or on Huggingface Spaces. 

* This tool is ideal for creators with many LORAs. One-click backup your entire portfolio.

* Keys are always stored locally on your browser. PRs welcome and appreciated.

* Usage: easy as 1️⃣->2️⃣->3️⃣, just follow the step numbers.



"""

with gr.Blocks(fill_width=True, css=css, delete_cache=(3600, 3600)) as demo:
    gr.Markdown(header_md, elem_classes="title")
    
    state = gr.State(value={})
    
    # Create browser state for API keys with error handling
    api_keys_state = gr.BrowserState(["", ""], storage_key="civitai2hf_api_keys")
    
    with gr.Row():
        show_btn = gr.Button("Show Modal", visible=False)
    with gr.Tabs() as tabs:
        with gr.TabItem("Search & Download"):
            with gr.Row():
                # Left column for filters (1/3 width)
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("## 2️⃣ Search Filters", container=True)
                        with gr.Accordion("Type & Model", open=True):
                            search_civitai_type = gr.CheckboxGroup(label="Type", choices=TYPE, value=["Checkpoint", "LORA"])
                            search_civitai_basemodel = gr.CheckboxGroup(label="Base Model", choices=BASEMODEL, value=[])
                            search_civitai_filetype = gr.CheckboxGroup(label="File type", visible=False,  choices=FILETYPE, value=["Model"])
                        with gr.Accordion("Search Options", open=True):
                            search_civitai_sort = gr.Radio(label="Sort", choices=SORT, value=SORT[0])
                            search_civitai_period = gr.Radio(label="Period", choices=PERIOD, value="Month")
                            search_civitai_limit =  gr.Slider(0, 100, value=50, step=10, label="Limit", info="Maximum items to query per page via API.")
                            search_civitai_page = gr.Slider(0, 10, 1, step=1, label="Num Pages", info="If 0, retrieve all pages")

                # Right column for results and download (2/3 width)
                with gr.Column(scale=2):
                    with gr.Group():
                        gr.Markdown("## 3️⃣ Search Query", container=True)
                        with gr.Row():
                            search_civitai_query = gr.Textbox(label="Query", placeholder="wan", lines=1)
                            search_civitai_tag = gr.Dropdown(label="Tag", choices=get_civitai_tag(), value=get_civitai_tag()[0], allow_custom_value=True)
                            search_civitai_user = gr.Textbox(label="Username", lines=1)
                        search_civitai_submit = gr.Button("Search Civitai API", variant="primary")
                    
                    with gr.Group():
                        gr.Markdown("## 4️⃣ Select Models to Backup from Search Results", container=True)
                        with gr.Row():
                            search_civitai_desc = gr.Markdown(value="", visible=False, elem_classes="desc")
                            search_civitai_json = gr.JSON(value={}, visible=False)
                        with gr.Row(equal_height=True):
                            with gr.Column(scale=9):
                                with gr.Accordion("👇 Gallery View ⚠️ Surprise NSFW + RAM Warning ⚠️", open=False):
                                    search_civitai_gallery = gr.Gallery(
                                        [], 
                                        label="Select from Results", 
                                        allow_preview=False, 
                                        columns=5, 
                                        elem_id="gallery", 
                                        object_fit="contain", 
                                        show_share_button=False, 
                                        interactive=True,
                                        preview=False,
                                        height="auto"
                                    )
                                with gr.Accordion("👇 List View", open=False):
                                    search_civitai_result_checkbox = gr.CheckboxGroup(label="", choices=[], value=[])
                                search_civitai_result = gr.Dropdown(label="Selected Models from Search", choices=[("", "")], value=[],
                                                                    allow_custom_value=True, visible=True, multiselect=True)
                                search_civitai_result_info = gr.Markdown("Standing by.", elem_classes="info")
                        with gr.Row():
                            search_civitai_add = gr.Button("5️⃣ Add Selected Models to DL List", variant="primary")
                            search_civitai_select_all = gr.Button("Select All", variant="secondary")
                    
                    with gr.Group():
                        gr.Markdown("## 6️⃣🚨 Ensure you have set your correct API key settings before moving on.", container=True)
                        dl_url = gr.Textbox(label="Download URL(s)", placeholder="https://civitai.com/api/download/models/28907\n...", value="", lines=3, max_lines=255, interactive=True)
                        with gr.Row():
                            newrepo_id = gr.Textbox(label="Your Repo ID", placeholder="yourid/yourrepo", value="", max_lines=1)
                            newrepo_type = gr.Radio(label="Repo Type", choices=["model", "dataset"], value="model")
                            with gr.Group():
                                is_private = gr.Checkbox(label="Private Repo", value=False)
                                is_rename = gr.Checkbox(label="Auto rename", value=True)
                        with gr.Row():
                            is_info = gr.Checkbox(label="Upload Civitai Metadata Files to HF", value=True)
                        
                        run_button = gr.Button(value="7️⃣ Create Repo & Backup Models", variant="primary")
                        uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=[]) # hidden
                        urls_md = gr.Markdown("<br><br>", elem_classes="result")
                        urls_remain = gr.Textbox("Remaining URLs", value="", show_copy_button=True, visible=False)

        with gr.TabItem("Settings 1️⃣"):
            with gr.Column():
                with gr.Group():
                    with gr.Row():
                        with gr.Column():
                            civitai_key = gr.Textbox(label="Your Civitai Key", value="", max_lines=1)
                            gr.Markdown("Your Civitai API key is available at [https://civitai.com/user/account](https://civitai.com/user/account).", elem_classes="info")
                        with gr.Column():
                            hf_token = gr.Textbox(label="Your HF write token", placeholder="hf_...", value="", max_lines=1)
                            gr.Markdown("Your token is available at [hf.co/settings/tokens](https://huggingface.co/settings/tokens).", elem_classes="info")
                    
                    # Add save status message
                    saved_message = gr.Markdown("✅ API keys saved to browser storage", visible=False)
                    
                    # Load saved keys when the app starts with error handling
                    @demo.load(inputs=[api_keys_state], outputs=[civitai_key, hf_token])
                    def load_api_keys(saved_keys):
                        try:
                            if not saved_keys or not isinstance(saved_keys, list) or len(saved_keys) != 2:
                                return "", ""
                            # Ensure we're not trying to use None values
                            civitai = saved_keys[0] or ""
                            hf = saved_keys[1] or ""
                            return civitai, hf
                        except Exception as e:
                            print(f"Error loading API keys: {e}")
                            return "", ""
                    
                    # Save keys when they change with improved error handling
                    @gr.on([civitai_key.change, hf_token.change], inputs=[civitai_key, hf_token], outputs=[api_keys_state, saved_message])
                    def save_api_keys(civitai_key, hf_token):
                        try:
                            # Ensure we're storing strings, not None
                            civitai = str(civitai_key or "")
                            hf = str(hf_token or "")
                            timestamp = time.strftime("%I:%M:%S %p")
                            
                            if civitai or hf:
                                return [civitai, hf], gr.Markdown(
                                    f"✅ API keys saved to browser storage at {timestamp}",
                                    visible=True
                                )
                            return ["", ""], gr.Markdown("", visible=False)
                        except Exception as e:
                            print(f"Error saving API keys: {e}")
                            return ["", ""], gr.Markdown("Error saving API keys", visible=True)
        with gr.TabItem("Known Issues"):
            with gr.Column():
                gr.Markdown("* Local browser key storage logic sometimes fails.\n* JSON Metadata for models fail to import.\n* WAN models seem to have a mind of its own. For the time being, it's recommended when searching for WAN models to UNCHECK `Wan Video` from filter and either use model tags or query directly.")
    with Modal(visible=True, elem_id="modal-window") as modal:
        gr.Markdown(common_sense)
    gr.on(
        triggers=[run_button.click],
        fn=download_civitai,
        inputs=[dl_url, civitai_key, hf_token, uploaded_urls, newrepo_id, newrepo_type, is_private, is_info, is_rename],
        outputs=[uploaded_urls, urls_md, urls_remain],
        queue=True,
    )
    gr.on(
        triggers=[search_civitai_submit.click, search_civitai_query.submit, search_civitai_user.submit],
        fn=search_civitai,
        inputs=[search_civitai_query, search_civitai_type, search_civitai_basemodel, search_civitai_sort,
                search_civitai_period, search_civitai_tag, search_civitai_user, search_civitai_limit,
                search_civitai_page, search_civitai_filetype, civitai_key, search_civitai_gallery, state],
        outputs=[search_civitai_result, search_civitai_desc, search_civitai_submit, search_civitai_query, search_civitai_gallery,
                 search_civitai_result_checkbox, search_civitai_result_info, state],
        queue=False,
        show_api=False,
    )
    search_civitai_result.change(select_civitai_item, [search_civitai_result, state], [search_civitai_desc, search_civitai_json, state], queue=False, show_api=False)\
    .success(update_civitai_checkbox, [search_civitai_result], [search_civitai_result_checkbox], queue=True, show_api=False)
    search_civitai_result_checkbox.select(from_civitai_checkbox, [search_civitai_result_checkbox], [search_civitai_result], queue=False, show_api=False)
    search_civitai_add.click(add_civitai_item, [search_civitai_result, dl_url], [dl_url], queue=False, show_api=False)
    search_civitai_select_all.click(select_civitai_all_item, [search_civitai_select_all, state], [search_civitai_select_all, search_civitai_result], queue=False, show_api=False)
    search_civitai_gallery.select(update_civitai_selection, [search_civitai_result, state], [search_civitai_result], queue=False, show_api=False)
    
    show_btn.click(lambda: Modal(visible=True), None, modal)

demo.queue()
demo.launch(ssr_mode=False)