import gradio as gr
from pygltflib import GLTF2, Texture, Image
import base64
import io

def modify_texture(glb_file, new_texture):
    # Load the GLB file
    gltf = GLTF2().load()

    # Convert the new texture image to use in GLB
    img_byte_arr = io.BytesIO()
    new_texture.save(img_byte_arr, format='PNG')
    encoded_img = base64.b64encode(img_byte_arr.getvalue()).decode('ascii')

    # Assuming there's at least one image in the original GLB
    if gltf.images:
        # Replace the first image with the new texture
        gltf.images[0].uri = f"data:image/png;base64,{encoded_img}"
    else:
        # Add new image if none exists
        new_image = Image(uri=f"data:image/png;base64,{encoded_img}")
        gltf.images.append(new_image)
        # Update texture to point to the new image
        new_texture = Texture(source=len(gltf.images)-1)
        gltf.textures.append(new_texture)

    # Save the modified GLB to a temporary file and return it
    output_file = io.BytesIO()
    gltf.save(output_file)
    output_file.seek(0)
    return output_file

with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            glb_input = gr.Model3D(label="Upload GLB File")
            texture_input = gr.Image(label="Upload New Texture Image")
        with gr.Column():
            modified_glb_output = gr.File(label="Download Modified GLB File")

    glb_input.change(modify_texture, inputs=[glb_input, texture_input], outputs=[modified_glb_output])

app.launch()