minor update (random seed etc)

app.py

@@ -27,27 +27,30 @@ with gr.Blocks(theme=theme) as demo:
             with gr.Tab("3D View (Slow)"):
                 generate_3d_button = gr.Button("Generate Terrain", variant="primary")
                 model_3d_output = gr.Model3D(
-                    camera_position=[90, 135, 512]
+                    camera_position=[90, 135, 512],
+                    clear_color=[0.0, 0.0, 0.0, 0.0],
+                    #display_mode = 'point_cloud'
                 )
 
         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="(3D Only) Crop Size")
+            random_seed = gr.Checkbox(value=True, label="Random Seed")
+            crop_size_slider = gr.Slider(minimum=128, maximum=768, step=64, value=768, label="(3D Only) Crop Size")
             vertex_count_slider = gr.Slider(minimum=0, maximum=10000, step=0, value=0, label="(3D Only) Vertex Count (Default: 0 - no reduction)")
             prefix_textbox = gr.Textbox(label="Prompt Prefix", value="A Sentinel-2 image of ")
             
         generate_2d_button.click(
             fn=generate_2d_view_output,
-            inputs=[prompt_input, num_inference_steps_slider, guidance_scale_slider, seed_number, prefix_textbox],
+            inputs=[prompt_input, num_inference_steps_slider, guidance_scale_slider, seed_number, random_seed, prefix_textbox],
             outputs=[rgb_output, elevation_output],
         )
 
         generate_3d_button.click(
             fn=generate_3d_view_output,
-            inputs=[prompt_input, num_inference_steps_slider, guidance_scale_slider, seed_number, crop_size_slider, vertex_count_slider, prefix_textbox],
-            outputs=[model_3d_output],
+            inputs=[prompt_input, num_inference_steps_slider, guidance_scale_slider, seed_number, random_seed, crop_size_slider, vertex_count_slider, prefix_textbox],
+            outputs=[rgb_output, elevation_output, model_3d_output],
         )
         
-        demo.queue().launch(share=True)
+        demo.queue().launch()

src/utils.py

@@ -10,13 +10,17 @@ import spaces
 pipe = build_pipe()
 
 @spaces.GPU
-def generate_terrain(prompt, num_inference_steps, guidance_scale, seed, prefix, crop_size=None):
+def generate_terrain(prompt, num_inference_steps, guidance_scale, seed, random_seed, prefix, crop_size=None):
     """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)
+    if random_seed:
+        generator = torch.Generator("cuda")
+    else:
+        generator = torch.Generator("cuda").manual_seed(seed)
+        
     image, dem = pipe(full_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator)
 
     if crop_size is not None:
@@ -50,6 +54,7 @@ def create_3d_mesh(rgb, elevation, n_clusters=1000):
     if n_clusters <= 0:
         # Generate full mesh without clustering
         vertices = points_3d
+        
         try:
             tri = Delaunay(points_2d)
             faces = tri.simplices
@@ -97,18 +102,29 @@ def create_3d_mesh(rgb, elevation, n_clusters=1000):
         mesh = trimesh.Trimesh(vertices=simplified_vertices, faces=valid_faces, vertex_colors=vertex_colors)
         return mesh
 
-def generate_3d_view_output(prompt, num_inference_steps, guidance_scale, seed, crop_size, vertex_count, prefix):
-    rgb, elevation = generate_terrain(prompt, num_inference_steps, guidance_scale, seed, prefix, crop_size)
+def create_3d_point_cloud(rgb, elevation):
+    height, width = elevation.shape
+    x, y = np.meshgrid(np.arange(width), np.arange(height))
+    points = np.stack([x.flatten(), y.flatten(), elevation.flatten()], axis=-1)
+    colors = rgb.reshape(-1, 3)
+
+    return trimesh.PointCloud(vertices=points, colors=colors)
+
+def generate_3d_view_output(prompt, num_inference_steps, guidance_scale, seed, random_seed, crop_size, vertex_count, prefix):
+    rgb, elevation = generate_terrain(prompt, num_inference_steps, guidance_scale, seed, random_seed, prefix, crop_size)
     
     mesh = create_3d_mesh(rgb, 500*elevation, n_clusters=vertex_count)
-
     with tempfile.NamedTemporaryFile(suffix=".obj", delete=False) as temp_file:
         mesh.export(temp_file.name)
         file_path = temp_file.name
+    # pc = create_3d_point_cloud(rgb, 500*elevation)
+    # with tempfile.NamedTemporaryFile(suffix=".ply", delete=False) as temp_file:
+    #     pc.export(temp_file.name, file_type="ply")
+    #     file_path = temp_file.name
 
-    return file_path
+    return rgb, elevation, file_path
 
-def generate_2d_view_output(prompt, num_inference_steps, guidance_scale, seed, prefix):
-    rgb, elevation = generate_terrain(prompt, num_inference_steps, guidance_scale, seed, prefix)
+def generate_2d_view_output(prompt, num_inference_steps, guidance_scale, seed, random_seed, prefix):
+    rgb, elevation = generate_terrain(prompt, num_inference_steps, guidance_scale, seed, random_seed, prefix)
 
     return rgb, elevation