Create app.py

app.py

@@ -0,0 +1,107 @@
+!pip install "huggingface_hub[hf_transfer]"
+!pip install -U "huggingface_hub[cli]"
+!pip install gradio trimesh scipy
+!HF_HUB_ENABLE_HF_TRANSFER=1
+!git clone https://github.com/PaulBorneP/MESA.git
+!cd MESA
+!mkdir weights
+!huggingface-cli download NewtNewt/MESA --local-dir weights
+
+import torch
+from MESA.pipeline_terrain import TerrainDiffusionPipeline
+import sys
+import gradio as gr
+import numpy as np
+import trimesh
+import tempfile
+import torch
+from scipy.spatial import Delaunay
+
+sys.path.append('MESA/')
+
+pipe = TerrainDiffusionPipeline.from_pretrained("./weights", torch_dtype=torch.float16)
+pipe.to("cuda")
+
+def generate_terrain(prompt, num_inference_steps, guidance_scale, seed, crop_size, prefix):
+    """Generates terrain data (RGB and elevation) from a text prompt."""
+    if prefix and not prefix.endswith(' '):
+        prefix += ' '  # Ensure prefix ends with a space
+
+    full_prompt = prefix + prompt
+    generator = torch.Generator("cuda").manual_seed(seed)
+    image, dem = pipe(full_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator)
+
+    # Center crop the image and dem
+    h, w, c = image[0].shape
+    start_h = (h - crop_size) // 2
+    start_w = (w - crop_size) // 2
+    end_h = start_h + crop_size
+    end_w = start_w + crop_size
+
+    cropped_image = image[0][start_h:end_h, start_w:end_w, :]
+    cropped_dem = dem[0][start_h:end_h, start_w:end_w, :]
+
+    return (255 * cropped_image).astype(np.uint8), 500*cropped_dem.mean(-1)
+
+def create_3d_mesh(rgb, elevation):
+    """Creates a 3D mesh from RGB and elevation data."""
+    x, y = np.meshgrid(np.arange(elevation.shape[1]), np.arange(elevation.shape[0]))
+    points = np.stack([x.flatten(), y.flatten()], axis=-1)
+    tri = Delaunay(points)
+
+    vertices = np.stack([x.flatten(), y.flatten(), elevation.flatten()], axis=-1)
+    faces = tri.simplices
+
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, vertex_colors=rgb.reshape(-1, 3))
+
+    return mesh
+
+def generate_and_display(prompt, num_inference_steps, guidance_scale, seed, crop_size, prefix):
+    """Generates terrain and displays it as a 3D model."""
+    rgb, elevation = generate_terrain(prompt, num_inference_steps, guidance_scale, seed, crop_size, prefix)
+    mesh = create_3d_mesh(rgb, elevation)
+
+    with tempfile.NamedTemporaryFile(suffix=".obj", delete=False) as temp_file:
+        mesh.export(temp_file.name)
+        file_path = temp_file.name
+
+    return file_path
+
+theme = gr.themes.Soft(primary_hue="red", secondary_hue="red", font=['arial'])
+
+with gr.Blocks(theme=theme) as demo:
+    with gr.Column(elem_classes="header"):
+        gr.Markdown("# MESA: Text-Driven Terrain Generation Using Latent Diffusion and Global Copernicus Data")
+        gr.Markdown("### Paul Borne–Pons, Mikolaj Czerkawski, Rosalie Martin, Romain Rouffet")
+        gr.Markdown('[[GitHub](https://github.com/PaulBorneP/MESA)] [[Model](https://huggingface.co/NewtNewt/MESA)] [[Dataset](https://huggingface.co/datasets/Major-TOM/Core-DEM)]')
+        
+    # Abstract Section
+    with gr.Column(elem_classes="abstract"):
+        gr.Markdown("MESA is a novel generative model based on latent denoising diffusion capable of generating 2.5D representations of terrain based on the text prompt conditioning supplied via natural language. The model produces two co-registered modalities of optical and depth maps.") # Replace with your abstract text
+        gr.Markdown("This is a test version of the demo app. Please be aware that MESA supports primarily complex, mountainous terrains as opposed to flat land")
+        gr.Markdown("The generated image is quite large, so for the full resolution (768) it might take a while to load the surface")
+
+    with gr.Row():
+        prompt_input = gr.Textbox(lines=2, placeholder="Enter a terrain description...")
+        generate_button = gr.Button("Generate Terrain", variant="primary")
+
+    model_output = gr.Model3D(
+        camera_position=[90, 180, 512]
+    )
+
+    with gr.Accordion("Advanced Options", open=False) as advanced_options:
+        num_inference_steps_slider = gr.Slider(minimum=10, maximum=1000, step=10, value=50, label="Inference Steps")
+        guidance_scale_slider = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, value=7.5, label="Guidance Scale")
+        seed_number = gr.Number(value=6378, label="Seed")
+        crop_size_slider = gr.Slider(minimum=128, maximum=768, step=64, value=512, label="Crop Size")
+        prefix_textbox = gr.Textbox(label="Prompt Prefix", value="A Sentinel-2 image of ")
+
+    generate_button.click(
+        fn=generate_and_display,
+        inputs=[prompt_input, num_inference_steps_slider, guidance_scale_slider, seed_number, crop_size_slider, prefix_textbox],
+        outputs=model_output,
+    )
+
+if __name__ == "__main__":
+    demo.launch(debug=True,
+                share=True)