app.py

import gradio as gr
import os
import random
import modules.constants as constants
import modules.version_info as version_info
import modules.storage as storage


user_dir = constants.TMPDIR
default_folder = "saved_models/3d_model_" + format(random.randint(1, 999999), "06d")

def getVersions():
    #return html_versions
    return version_info.versions_html()

def load_data(query_params, model_3d, image_slider):
    # set default values or pull from querystring
    model_url = query_params.get("3d", None) if query_params else None
    hm_url = query_params.get("hm", None) if query_params else None
    img_url = query_params.get("image", None) if query_params else None
    slider_images = []
    if hm_url:
        slider_images.append(hm_url)
    if img_url:
        slider_images.append(img_url)
    if not slider_images:
        slider_images = ["images/beeuty_545jlbh1_v12_alpha96_300dpi.png", "images/beeuty_545jlbh1_v12_alpha96_300dpi_depth.png"]
    if not model_url:
        model_url = "models/beeuty_545jlbh1_300dpi.glb"
    return model_url, slider_images

def process_upload(files, current_model, current_images):
    """
    Process uploaded files and assign them to the appropriate component based on file extension.
    
    Files with extensions in [".glb", ".gltf", ".obj", ".ply"] are sent to the Model3D component.
    Files with extensions in [".png", ".jpg", ".jpeg"] are sent to the ImageSlider component.
    
    The function merges the uploaded files with current data. If a file for a component is not 
    provided in the upload (i.e. not exactly 1 model file or not exactly 2 image files), then the
    original data will be retained for that component. If an upload is provided, it will replace 
    the corresponding value.
    
    For the ImageSlider, if a single image is provided in the upload, it will update only the first
    image slot, leaving the second slot unchanged.
    """
    extracted_model = None
    extracted_images = []
    
    # Ensure files is a list.
    if not isinstance(files, list):
        files = [files]
        
    for f in files:
        # f can be a file path (string) or an object with attribute `name`
        file_name = f.name if hasattr(f, "name") else f
        ext = os.path.splitext(file_name)[1].lower()
        
        if ext in constants.model_extensions:
            if extracted_model is None:
                extracted_model = file_name
        elif ext in constants.image_extensions:
            if len(extracted_images) < 2:
                extracted_images.append(file_name)
                
    # Merge results with current data.
    updated_model = extracted_model if extracted_model is not None else current_model
    
    # Convert current_images if it's a tuple or a single item.
    if isinstance(current_images, tuple):
        current_images = list(current_images)
    elif current_images is not None and not isinstance(current_images, list):
        current_images = [current_images]
    
    # For the image slider, we expect a list of exactly 2 images.
    # Start with current images (or use defaults if None).
    if current_images is None or not isinstance(current_images, list):
        new_images = [None, None]
    else:
        new_images = current_images + [None] * (2 - len(current_images))
        new_images = new_images[:2]
    
    # If at least one image is uploaded, update the corresponding slot(s).
    for i in range(len(extracted_images)):
        if i < 2:
            new_images[i] = extracted_images[i]
    
    return updated_model, new_images

gr.set_static_paths(paths=["images/", "models/", "assets/"])
with gr.Blocks(css_paths="style_20250503.css", title="3D viewer", theme='Surn/beeuty',delete_cache=(21600,86400), fill_width=True) as viewer3d:
    gr.Markdown("# 3D Model Viewer")
    
    with gr.Row():
        with gr.Column():
            model_3d = gr.Model3D(
                label="3D Model",
                value=None,
                height=480,
                elem_id="model_3d", key="model_3d", clear_color=[1.0, 1.0, 1.0, 0.1], 
                elem_classes="centered solid imgcontainer", interactive=True
                
            )
            image_slider = gr.ImageSlider(
                label="2D Images",
                value=None,
                height=480,
                elem_id="image_slider", key="image_slider",
                type="filepath"
            )
        
    with gr.Row():
        gr.Markdown("## Upload your own files")
        gr.Markdown("### Supported formats: " + ", ".join([f"`{ext}`" for ext in constants.upload_file_types]))
    with gr.Row():
        upload_btn = gr.UploadButton(
            "Upload 3D Files", elem_id="upload_btn", key="upload_btn",
            file_count="multiple",
            file_types=constants.upload_file_types
        )
        
    with gr.Row():
        # New textbox for folder name.
        folder_name_box = gr.Textbox(
            label="Folder Name",
            value=default_folder,
            elem_id="folder_name",
            key="folder_name",
            placeholder="Enter folder name...",
            elem_classes="solid centered"
        )
        permalink_button = gr.Button("Generate Permalink", elem_id="permalink_button", key="permalink_button", elem_classes="solid small centered")

    with gr.Row(visible=False, elem_id="permalink_row") as permalink_row:
        permalink = gr.Textbox(
            show_copy_button=True,
            label="Permalink",
            elem_id="permalink",
            key="permalink",
            elem_classes="solid centered",
            max_lines=5,
            lines=3
        )
        gr.Markdown("### Copy the link above to share your model and images.", elem_classes="solid centered",)
    with gr.Row():
        gr.HTML(value=getVersions(), visible=True, elem_id="versions")
    
    # Use JavaScript to pass the query parameters to your callback.
    viewer3d.load(
        load_data,
        inputs=[gr.JSON(), model_3d, image_slider],
        outputs=[model_3d, image_slider],
        js="""() => {
            const params = Object.fromEntries(new URLSearchParams(window.location.search));
            return params;
        }""",
        scroll_to_output=True
    )

    # Process uploaded files to update the Model3D or ImageSlider component.
    upload_btn.upload(
        process_upload, 
        inputs=[upload_btn, model_3d, image_slider],
        outputs=[model_3d, image_slider],
        scroll_to_output=True,
        api_name="process_upload",
        show_progress=True

    )
    # Generate a permalink based on the current model, images, and folder name.
    permalink_button.click(
        lambda model, images, folder: storage.upload_files_to_repo(
            files=[model] + list(images),
            repo_id="Surn/Storage",
            folder_name=folder,
            create_permalink=True,
            repo_type="dataset"
        )[1],  # Extract the permalink from the returned tuple if criteria met.
        inputs=[model_3d, image_slider, folder_name_box],
        outputs=[permalink],
        scroll_to_output=True
    ).then(
        lambda link: gr.update(visible=True) if link and len(link) > 0 else gr.update(visible=False),
        inputs=[permalink],
        outputs=[permalink_row]
    )

if __name__ == "__main__":
    viewer3d.launch(
        allowed_paths=["assets", "assets/", "./assets", "images/", "./images", 'e:/TMP', 'models/', '3d_model_viewer/'],
        favicon_path="./assets/favicon.ico", show_api=True, strict_cors=False
    )