Update app.py

app.py

@@ -6,7 +6,7 @@ import os
 import plyfile
 import open_clip
 
-# Load OpenCLIP for prompt conditioning (runs efficiently on CPU or GPU)
+# Load OpenCLIP for prompt conditioning
 model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79K')
 model.eval()
 tokenizer = open_clip.get_tokenizer('ViT-B-32')
@@ -14,7 +14,7 @@ tokenizer = open_clip.get_tokenizer('ViT-B-32')
 class PersistentCortex:
     def __init__(self, num_gaussians=8000):
         self.num = num_gaussians
-        # Initialize a compact spherical distribution (similar to your working PLY)
+        # Initialize compact spherical distribution
         angles = torch.rand(num_gaussians, 2) * 2 * np.pi
         radius = torch.rand(num_gaussians).pow(1/3) * 0.6
         self.positions = torch.stack([
@@ -24,35 +24,38 @@ class PersistentCortex:
         ], dim=1)
         
         self.scales = torch.exp(torch.randn(num_gaussians, 3) * -2.5 - 2.0)
-        self.colors = torch.rand(num_gaussians, 3) * 0.7 + 0.3  # neutral gray start
+        self.colors = torch.rand(num_gaussians, 3) * 0.7 + 0.3
         self.opacities = torch.sigmoid(torch.randn(num_gaussians) * 1.5 + 2.0)
         self.rotations = torch.nn.functional.normalize(torch.randn(num_gaussians, 4), dim=-1)
 
     def evolve_from_image(self, image: Image.Image, steps=800):
         img_array = np.array(image.resize((256, 256))) / 255.0
-        target = torch.tensor(img_array, dtype=torch.float32)
+        target = torch.tensor(img_array, dtype=torch.float32).permute(2, 0, 1)  # (3, 256, 256)
+
+        # Create sampling grid: scatter points across [-1,1] x [-1,1]
+        proj = torch.rand(self.num, 2) * 2 - 1  # (N, 2) uniform in view
 
         for _ in range(steps):
-            # Simple orthographic projection to 2D for guidance
-            proj = self.positions[:, :2].clone()
-            proj = proj / proj.abs().max()  # normalize to [-1, 1]
-            grid = proj.unsqueeze(0).unsqueeze(0)  # (1,1,N,2)
+            grid = proj.unsqueeze(0).unsqueeze(0)  # (1, 1, N, 2)
 
             sampled = torch.nn.functional.grid_sample(
-                target.permute(2, 0, 1).unsqueeze(0),
+                target.unsqueeze(0),
                 grid,
                 mode='bilinear',
                 padding_mode='border',
                 align_corners=True
-            ).squeeze(0).squeeze(0)  # (N, 3)
+            ).squeeze(0).squeeze(0).t()  # (N, 3)
 
-            # Attract positions toward brighter areas
             brightness = sampled.mean(dim=1)
-            attraction = (brightness - brightness.mean()) * 0.015
-            self.positions[:, :2] += attraction.unsqueeze(1) * proj
+            attraction_strength = (brightness - brightness.mean()) * 0.02
+            proj += attraction_strength.unsqueeze(1) * proj.normalized()
+
+            # Color adaptation
+            avg_color = target.mean(dim=[1,2])
+            self.colors = torch.lerp(self.colors, avg_color.repeat(self.num, 1), 0.02)
 
-            # Gradually adopt average color from target
-            self.colors = torch.lerp(self.colors, target.mean(dim=[0,1]).repeat(self.num, 1), 0.02)
+        # Map projected points back to 3D sphere surface (simple radial projection)
+        self.positions[:, :2] = proj * (self.positions[:, :2].norm(dim=1).unsqueeze(1) + 0.1)
 
         return self
 
@@ -64,7 +67,6 @@ class PersistentCortex:
             text_emb = model.encode_text(text_tokens).float()
             text_emb = text_emb / text_emb.norm(dim=-1, keepdim=True)
 
-        # Simple linear projection to RGB shift (deterministic for reproducibility)
         proj = torch.nn.Linear(512, 3, bias=False)
         torch.nn.init.normal_(proj.weight, std=0.2)
         color_shift = proj(text_emb)
@@ -73,17 +75,16 @@ class PersistentCortex:
         return self
 
     def export_ply(self, path="output.ply"):
-        # Prepare vertex data matching gsplat.js expected properties
         zeros = np.zeros((self.num, 3), dtype=np.float32)
-        log_scales = np.log(self.scales.cpu().numpy()).astype(np.float32)
+        log_scales = np.log(self.scales.cpu().numpy())
 
         vertex_data = np.core.records.fromarrays([
-            self.positions.cpu().numpy().astype(np.float32).T,
-            zeros.T,  # normals (unused)
-            self.colors.cpu().numpy().T,
-            self.opacities.cpu().numpy().reshape(1, -1).T,
-            log_scales.T,
-            self.rotations.cpu().numpy().T
+            self.positions.cpu().numpy(),
+            zeros,
+            self.colors.cpu().numpy(),
+            self.opacities.cpu().numpy()[:, np.newaxis],
+            log_scales,
+            self.rotations.cpu().numpy()
         ], names='x,y,z,nx,ny,nz,f_dc_0,f_dc_1,f_dc_2,opacity,scale_0,scale_1,scale_2,rot_0,rot_1,rot_2,rot_3')
 
         el = plyfile.PlyElement.describe(vertex_data, 'vertex')
@@ -92,7 +93,7 @@ class PersistentCortex:
 
 def process(image: Image.Image, prompt: str = ""):
     if image is None:
-        raise ValueError("Please upload an image to proceed.")
+        raise gr.Error("Please upload an image.")
     
     cortex = PersistentCortex(num_gaussians=8000)
     cortex.evolve_from_image(image, steps=800)
@@ -101,7 +102,6 @@ def process(image: Image.Image, prompt: str = ""):
 
     ply_path = cortex.export_ply("/tmp/output.ply")
 
-    # In Hugging Face Spaces, /tmp files are automatically served under /files/
     viewer_html = f"""
     <div id="viewer" style="width:100%; height:600px; background:#000;"></div>
     <script type="module">
@@ -121,9 +121,6 @@ def process(image: Image.Image, prompt: str = ""):
     controls.autoRotate = false;
     controls.enableDamping = true;
     controls.dampingFactor = 0.05;
-    controls.rotateSpeed = 1.0;
-    controls.zoomSpeed = 1.2;
-    controls.panSpeed = 0.8;
 
     await SPLAT.Loader.LoadAsync("/files/output.ply", scene);
 
@@ -136,34 +133,26 @@ def process(image: Image.Image, prompt: str = ""):
     </script>
     """
 
-    status = "3D Gaussian splat generated and evolved from your image"
+    status = f"Persistent 3D representation evolved from image"
     if prompt:
-        status += f" with prompt conditioning: '{prompt}'."
-    else:
-        status += "."
+        status += f" and conditioned on prompt: '{prompt}'"
+    status += "."
 
     return viewer_html, status
 
 with gr.Blocks(title="Persistent 3D Cortex Demo") as demo:
     gr.Markdown("# Persistent 3D Cortex – Interactive Demo")
-    gr.Markdown("""
-    Upload an image and optionally add a text prompt.  
-    The system evolves a persistent 3D Gaussian splat representation influenced by the image content and prompt.
-    """)
+    gr.Markdown("Upload an image and optionally add a text prompt to generate an evolving 3D Gaussian splat representation.")
     
     with gr.Row():
-        img_input = gr.Image(type="pil", label="Input Image")
-        prompt_input = gr.Textbox(label="Prompt (e.g., 'shiny red apple', 'futuristic city')", placeholder="Optional text prompt")
+        img_input = gr.Image(type="pil", label="Input Image (required)")
+        prompt_input = gr.Textbox(label="Prompt (e.g., 'glowing blue crystal')", placeholder="Optional descriptive text")
     
-    generate_btn = gr.Button("Generate & Evolve 3D", variant="primary")
+    generate_btn = gr.Button("Generate Persistent 3D", variant="primary")
     
-    viewer_output = gr.HTML(label="Interactive 3D Viewer")
+    viewer_output = gr.HTML()
     status_output = gr.Textbox(label="Status")
 
-    generate_btn.click(
-        fn=process,
-        inputs=[img_input, prompt_input],
-        outputs=[viewer_output, status_output]
-    )
+    generate_btn.click(process, inputs=[img_input, prompt_input], outputs=[viewer_output, status_output])
 
 demo.launch()