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()
    # Process URLs and download files if needed.

def process_url(url, default_ext=".png"):
    """Download file from URL if it's a remote URL and return its local path.
    
    Performs HuggingFace authentication if the URL requires it.
    The caller can pass an appropriate default_ext (e.g. ".glb" for models).
    Uses huggingface_hub library for HuggingFace URLs for better authentication.
    """
    if not url:
        return None
        
    # If it's already a local file, return it.
    if os.path.exists(url) or not (url.startswith('http://') or url.startswith('https://')):
        return url
    
    # Parse URL to get components
    try:
        import urllib.request
        from urllib.parse import urlparse
        
        # Create filename from URL.
        parsed_url = urlparse(url)
        filename = os.path.basename(parsed_url.path)
        if not filename:
            filename = f"downloaded_{hash(url) % 10000}.file"
            
        # Add extension if missing.
        ext = os.path.splitext(filename)[1].lower()
        if not ext:
            filename += default_ext
            
        # Create local path.
        local_path = os.path.join(constants.TMPDIR, filename)
        
        # If the file is hosted on HuggingFace, use huggingface_hub
        if 'huggingface.co' in url or 'hf.co' in url:
            try:
                from huggingface_hub import login, hf_hub_download

                # Log in to HuggingFace
                login(token=constants.HF_API_TOKEN)

                # Extract repo information from URL
                # Format: https://huggingface.co/datasets/{repo_id}/resolve/main/{path}
                if '/datasets/' in url and '/resolve/main/' in url:
                    parts = url.split('/datasets/')[1].split('/resolve/main/')
                    repo_id = parts[0]
                    
                    # The remaining path may contain subfolders and filename
                    full_path = parts[1]
                    
                    # Extract the filename and subfolder
                    if '/' in full_path:
                        subfolder, filename = full_path.rsplit('/', 1)
                    else:
                        subfolder = None
                        filename = full_path
                    
                    print(f"Downloading from HF repo '{repo_id}', filename '{filename}', subfolder '{subfolder}'")
                    
                    # Download using huggingface_hub
                    local_path = hf_hub_download(
                        repo_id=repo_id,
                        filename=filename,
                        subfolder=subfolder,
                        repo_type="dataset",
                        local_dir=constants.TMPDIR,
                        local_dir_use_symlinks=False
                    )
                    return local_path
                else:
                    # Fall back to standard download for other HF URLs
                    print("URL format not recognized for huggingface_hub download, falling back to standard method")
            except Exception as e:
                print(f"Error using huggingface_hub download: {e}, falling back to standard method")
        
        # Standard download for non-HF URLs or as fallback
        print(f"Downloading {url} to {local_path}")
        urllib.request.urlretrieve(url, local_path)
        return local_path
            
    except Exception as e:
        print(f"Error downloading file {url}: {e}")
        return url  # Return original URL if download fails


def load_data(request: gr.Request, model_3d, image_slider):
    """
    Load data from query parameters, download files if needed,
    and use current component values as defaults if no query parameters are provided.
    
    If query parameters are provided, generate a permalink using storage.generate_permalink_from_urls.
    
    Parameters:
        request: Gradio request object containing query parameters.
        model_3d: Current value or component for the 3D model.
        image_slider: Current value or component for the image slider.
    
    Returns:
        tuple: (model_url, slider_images, permalink)
               - model_url: processed URL for the 3D model.
               - slider_images: processed list of image URLs.
               - permalink: a generated permalink if query parameters were provided,
                 or an empty string if not.
    """
    # Parse query parameters.
    query_params = dict(request.query_params) if request is not None else {}
    
    # Extract URLs from query parameters.
    model_url = query_params.get("3d", None)
    hm_url = query_params.get("hm", None)
    img_url = query_params.get("image", None)

    if model_url is None and hm_url is None and img_url is None:
        # No URLs provided, return default values.
        query_params = {}
    
    # Process the model URL if provided.
    if model_url:
        model_url = process_url(model_url, default_ext=".glb")
    
    # Process image URLs if provided.
    slider_images = []
    if img_url:
        local_img = process_url(img_url, default_ext=".png")
        if local_img:
            slider_images.append(local_img)
    if hm_url:
        local_hm = process_url(hm_url, default_ext=".png")
        if local_hm:
            slider_images.append(local_hm)

    
    # Set default values if no URLs provided:
    default_model = getattr(model_3d, "value", model_3d)
    default_images = getattr(image_slider, "value", image_slider)
    
    if not slider_images:
        slider_images = default_images if not default_images == (None,None) else constants.default_slider_images
        
    if not model_url:
        model_url = default_model if default_model else constants.default_model_3d

    # If any query parameters were provided, generate a permalink.
    permalink = ""
    if query_params:
        try:
            # Use the helper function defined in storage.py
            permalink = storage.generate_permalink_from_urls(model_url, hm_url, img_url)
        except Exception as e:
            print(f"Error generating permalink: {e}")
    
    return model_url, slider_images, permalink

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,
                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="100%",
                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 permalink 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=[model_3d, image_slider],
        outputs=[model_3d, image_slider, permalink],
        scroll_to_output=True
    ).then(
        # If the returned permalink (link) is non-empty then make the permalink row visible
        # and disable the permalink button; otherwise, hide the row and enable the button.
        lambda link: (gr.update(visible=True), gr.update(interactive=False))
            if link and len(link) > 0
            else (gr.update(visible=False), gr.update(interactive=True)),
        inputs=[permalink],
        outputs=[permalink_row, permalink_button]
    )

    # 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
    ).then(
        # After a successful upload, enable the permalink button.
        lambda m, i: gr.update(interactive=True),
        inputs=[model_3d, image_slider],
        outputs=[permalink_button]
    )
    # 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
    )