Update app.py

app.py

@@ -2,95 +2,165 @@ import gradio as gr
 import torch
 import numpy as np
 from PIL import Image
-import tempfile
-import plyfile  # Add to requirements.txt
-
-class SimplePersistentCortex:
-    def __init__(self, num_gaussians=10000):
-        self.num_gaussians = num_gaussians
-        # Initialize clustered Gaussians
-        self.positions = torch.randn(num_gaussians, 3) * 0.4
-        self.scales = torch.exp(torch.randn(num_gaussians, 3) * 0.6 - 1.0)  # Positive scales
-        self.colors_dc = torch.rand(num_gaussians, 3) * 0.5 + 0.25  # Base colors (SH degree 0)
-        self.opacities = torch.sigmoid(torch.randn(num_gaussians) * 1.5 + 2.0)  # Mostly visible
+import os
+import plyfile
+import open_clip
+
+# Load OpenCLIP for prompt conditioning (runs efficiently on CPU or GPU)
+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')
+
+class PersistentCortex:
+    def __init__(self, num_gaussians=8000):
+        self.num = num_gaussians
+        # Initialize a compact spherical distribution (similar to your working PLY)
+        angles = torch.rand(num_gaussians, 2) * 2 * np.pi
+        radius = torch.rand(num_gaussians).pow(1/3) * 0.6
+        self.positions = torch.stack([
+            radius * torch.sin(angles[:, 0]) * torch.cos(angles[:, 1]),
+            radius * torch.sin(angles[:, 0]) * torch.sin(angles[:, 1]),
+            radius * torch.cos(angles[:, 0])
+        ], 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.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(self, target_image: Image.Image, steps: int = 100):
-        target_np = np.array(target_image.resize((128, 128))) / 255.0
-        target_tensor = torch.tensor(target_np, dtype=torch.float32).mean(dim=(0,1))  # Average color
+    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)
 
         for _ in range(steps):
-            # Attract positions to center with noise
-            self.positions += torch.randn_like(self.positions) * 0.01
-            self.positions *= 0.99  # Dampen spread
+            # 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)
+
+            sampled = torch.nn.functional.grid_sample(
+                target.permute(2, 0, 1).unsqueeze(0),
+                grid,
+                mode='bilinear',
+                padding_mode='border',
+                align_corners=True
+            ).squeeze(0).squeeze(0)  # (N, 3)
 
-            # Adapt colors toward target
-            self.colors_dc += (target_tensor - self.colors_dc.mean(dim=0)) * 0.02
-            self.colors_dc = torch.clamp(self.colors_dc, 0, 1)
+            # Attract positions toward brighter areas
+            brightness = sampled.mean(dim=1)
+            attraction = (brightness - brightness.mean()) * 0.015
+            self.positions[:, :2] += attraction.unsqueeze(1) * proj
+
+            # Gradually adopt average color from target
+            self.colors = torch.lerp(self.colors, target.mean(dim=[0,1]).repeat(self.num, 1), 0.02)
+
+        return self
 
-            # Vary scales for density
-            self.scales *= torch.exp(torch.randn_like(self.scales) * 0.03)
+    def condition_on_prompt(self, prompt: str):
+        if not prompt.strip():
+            return self
+        text_tokens = tokenizer([prompt])
+        with torch.no_grad():
+            text_emb = model.encode_text(text_tokens).float()
+            text_emb = text_emb / text_emb.norm(dim=-1, keepdim=True)
 
-            # Boost opacity
-            self.opacities = torch.clamp(self.opacities + 0.01, 0, 1)
+        # 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)
 
+        self.colors = torch.clamp(self.colors + color_shift.repeat(self.num, 1) * 0.3, 0, 1)
         return self
 
-    def export_ply(self, path: str):
-        positions = self.positions.cpu().numpy()
-        scales = np.log(self.scales.cpu().numpy() + 1e-9)
-        opacities = np.arctanh(self.opacities.cpu().numpy().clip(1e-6, 1-1e-6))
-        rot = self.rotations.cpu().numpy()
-        sh_dc = self.colors_dc.cpu().numpy().reshape(-1, 3)  # Simplified SH
-
-        # Create vertex data
-        vertex_data = np.empty(self.num_gaussians, dtype=[
-            ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
-            ('f_dc_0', 'f4'), ('f_dc_1', 'f4'), ('f_dc_2', 'f4'),
-            ('opacity', 'f4'),
-            ('scale_0', 'f4'), ('scale_1', 'f4'), ('scale_2', 'f4'),
-            ('rot_0', 'f4'), ('rot_1', 'f4'), ('rot_2', 'f4'), ('rot_3', 'f4')
-        ])
-
-        vertex_data['x'], vertex_data['y'], vertex_data['z'] = positions[:, 0], positions[:, 1], positions[:, 2]
-        vertex_data['f_dc_0'], vertex_data['f_dc_1'], vertex_data['f_dc_2'] = sh_dc[:, 0], sh_dc[:, 1], sh_dc[:, 2]
-        vertex_data['opacity'] = opacities
-        vertex_data['scale_0'], vertex_data['scale_1'], vertex_data['scale_2'] = scales[:, 0], scales[:, 1], scales[:, 2]
-        vertex_data['rot_0'], vertex_data['rot_1'], vertex_data['rot_2'], vertex_data['rot_3'] = rot[:, 0], rot[:, 1], rot[:, 2], rot[:, 3]
+    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)
+
+        vertex_data = np.core.records.fromarrays([
+            self.positions.cpu().numpy().astype(np.float32),
+            zeros,  # normals (unused)
+            self.colors.cpu().numpy().astype(np.float32),
+            self.opacities.cpu().numpy().astype(np.float32),
+            log_scales,
+            self.rotations.cpu().numpy().astype(np.float32)
+        ], 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')
-        plyfile.PlyData([el]).write(path)
+        plyfile.PlyData([el], text=True).write(path)
+        return os.path.abspath(path)
 
 def process(image: Image.Image, prompt: str = ""):
-    if image is None:
-        return None, "Error: Please upload an image."
-    
-    cortex = SimplePersistentCortex()
-    cortex.evolve(image)
-    
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".ply") as tmp:
-        cortex.export_ply(tmp.name)
-        ply_path = tmp.name
-    
-    status = "Persistent 3D Gaussian splat generated. Use mouse to orbit/zoom in viewer."
+    cortex = PersistentCortex(num_gaussians=8000)
+    cortex.evolve_from_image(image, steps=800)
     if prompt:
-        status += f" (Prompt: {prompt})"
-    
-    return ply_path, status
+        cortex.condition_on_prompt(prompt)
+
+    ply_path = cortex.export_ply("/tmp/output.ply")
 
-with gr.Blocks(title="Persistent 3D Cortex Demo v3") as demo:
-    gr.Markdown("# Persistent 3D Cortex Demo")
-    gr.Markdown("Upload an image to generate an evolving 3D Gaussian representation.")
+    # 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">
+    import * as SPLAT from "https://cdn.jsdelivr.net/npm/gsplat@latest";
+
+    const container = document.getElementById('viewer');
+    const canvas = document.createElement('canvas');
+    canvas.style.width = '100%';
+    canvas.style.height = '100%';
+    container.appendChild(canvas);
+
+    const scene = new SPLAT.Scene();
+    const camera = new SPLAT.Camera();
+    const renderer = new SPLAT.WebGLRenderer({{ canvas }});
+    const controls = new SPLAT.OrbitControls(camera, canvas);
+
+    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);
+
+    function animate() {{
+        controls.update();
+        renderer.render(scene, camera);
+        requestAnimationFrame(animate);
+    }}
+    animate();
+    </script>
+    """
+
+    status = "3D Gaussian splat generated and evolved from your image"
+    if prompt:
+        status += f" with prompt conditioning: '{prompt}'."
+    else:
+        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.
+    """)
     
     with gr.Row():
         img_input = gr.Image(type="pil", label="Input Image")
-        prompt_input = gr.Textbox(label="Prompt (optional)")
+        prompt_input = gr.Textbox(label="Prompt (e.g., 'shiny red apple', 'futuristic city')", placeholder="Optional text prompt")
     
-    btn = gr.Button("Generate 3D")
+    generate_btn = gr.Button("Generate & Evolve 3D", variant="primary")
     
-    model_output = gr.Model3D(label="Interactive Gaussian Splat Viewer", clear_color=[0.0, 0.0, 0.0, 1.0])
-    text_output = gr.Textbox(label="Status")
-    
-    btn.click(process, inputs=[img_input, prompt_input], outputs=[model_output, text_output])
+    viewer_output = gr.HTML(label="Interactive 3D Viewer")
+    status_output = gr.Textbox(label="Status")
+
+    generate_btn.click(
+        fn=process,
+        inputs=[img_input, prompt_input],
+        outputs=[viewer_output, status_output]
+    )
 
 demo.launch()