custom endpoints

app.py

@@ -12,10 +12,15 @@ from PIL import Image
 from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
+import requests
+import base64
+import io
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(TMP_DIR, exist_ok=True)
 
+NODE_SERVER_UPLOAD_URL = "https://viverse-backend.onrender.com/api/upload-rigged-model"
+
 # Funciones auxiliares
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
@@ -120,6 +125,153 @@ def extract_glb(
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
+@spaces.GPU(duration=180)
+def generate_model_from_images_and_upload(
+    image_inputs: List[str], 
+    input_type: str, 
+    seed_val: int,
+    ss_guidance_strength_val: float,
+    ss_sampling_steps_val: int,
+    slat_guidance_strength_val: float,
+    slat_sampling_steps_val: int,
+    multiimage_algo_val: str,
+    mesh_simplify_val: float,
+    texture_size_val: int,
+    model_description: str,
+    req: gr.Request
+) -> str:
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+    
+    pil_images = []
+    image_basenames = []
+    print(f"Received image_inputs: {image_inputs}, input_type: {input_type}")
+
+    for i, img_data in enumerate(image_inputs):
+        try:
+            print(f"Processing image {i+1}/{len(image_inputs)} with type '{input_type}'")
+            if input_type == "url":
+                print(f"Fetching image from URL: {img_data}")
+                response_img = requests.get(img_data, stream=True, timeout=30)
+                response_img.raise_for_status()
+                img = Image.open(response_img.raw)
+                image_basenames.append(os.path.basename(img_data).split('.')[0] or f"image_{i}")
+            elif input_type == "base64":
+                print(f"Decoding base64 image data (first 30 chars): {img_data[:30]}...")
+                # Ensure correct padding for base64
+                missing_padding = len(img_data) % 4
+                if missing_padding:
+                    img_data += '=' * (4 - missing_padding)
+                img_bytes = base64.b64decode(img_data)
+                img = Image.open(io.BytesIO(img_bytes))
+                image_basenames.append(f"base64_image_{i}")
+            elif input_type == "filepath": 
+                print(f"Opening image from filepath: {img_data}")
+                img = Image.open(img_data)
+                image_basenames.append(os.path.basename(img_data).split('.')[0] or f"image_{i}")
+            else:
+                print(f"Unsupported input_type: {input_type}")
+                raise ValueError(f"Unsupported input_type: {input_type}")
+            
+            print(f"Image {i+1} loaded, mode: {img.mode}, size: {img.size}. Preprocessing...")
+            # Ensure image is in RGB format if it's not, e.g. RGBA or P
+            if img.mode == 'RGBA' or img.mode == 'P':
+                print(f"Converting image {i+1} from {img.mode} to RGB")
+                img = img.convert('RGB')
+
+            processed_img = pipeline.preprocess_image(img)
+            pil_images.append(processed_img)
+            print(f"Image {i+1} processed and added.")
+
+        except Exception as e:
+            print(f"Error processing image {i} ('{str(img_data)[:50]}...'): {e}")
+            import traceback
+            traceback.print_exc()
+            raise gr.Error(f"Failed to load or process input image {i} ({input_type}): {e}")
+
+    if not pil_images:
+        print("No valid images could be processed.")
+        raise gr.Error("No valid images could be processed.")
+    
+    print(f"Total PIL images for pipeline: {len(pil_images)}")
+
+    print("Running multi-image pipeline...")
+    outputs = pipeline.run_multi_image(
+        pil_images,
+        seed=seed_val,
+        formats=["gaussian", "mesh"],
+        preprocess_image=False, 
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps_val,
+            "cfg_strength": ss_guidance_strength_val,
+        },
+        slat_sampler_params={
+            "steps": slat_sampling_steps_val,
+            "cfg_strength": slat_guidance_strength_val,
+        },
+        mode=multiimage_algo_val,
+    )
+    print("Multi-image pipeline completed.")
+    
+    gs_result = outputs['gaussian'][0]
+    mesh_result = outputs['mesh'][0]
+
+    print(f"Extracting GLB with simplify: {mesh_simplify_val}, texture_size: {texture_size_val}")
+    glb_data = postprocessing_utils.to_glb(gs_result, mesh_result, simplify=mesh_simplify_val, texture_size=texture_size_val, verbose=False)
+    
+    temp_glb_filename = 'temp_output_image_model.glb'
+    temp_glb_path = os.path.join(user_dir, temp_glb_filename)
+    print(f"Exporting GLB to temporary path: {temp_glb_path}")
+    glb_data.export(temp_glb_path)
+    
+    torch.cuda.empty_cache()
+    print("CUDA cache cleared.")
+
+    print(f"Uploading GLB from {temp_glb_path} to {NODE_SERVER_UPLOAD_URL}")
+    persistent_url = None
+    upload_prompt_name = model_description or "_".join(filter(None, image_basenames)) or "imagen_generated_model"
+    # Sanitize upload_prompt_name further for safety
+    upload_prompt_name = "".join(c if c.isalnum() or c in ['_', '-'] else '_' for c in upload_prompt_name)[:50]
+
+
+    try:
+        with open(temp_glb_path, "rb") as f:
+            files = {"modelFile": (temp_glb_filename, f, "model/gltf-binary")}
+            payload = {
+                "clientType": "playcanvas", 
+                "prompt": upload_prompt_name, 
+                "modelStage": "imagen_trellis_tpose" 
+            }
+            print(f"Upload payload to Node.js: {payload}")
+            response = requests.post(NODE_SERVER_UPLOAD_URL, files=files, data=payload, timeout=120)
+            response.raise_for_status()
+            result = response.json()
+            persistent_url = result.get("persistentUrl")
+            if not persistent_url:
+                print(f"No persistent URL in Node.js server response: {result}")
+                raise ValueError("Upload successful, but no persistent URL returned from Node.js server")
+            print(f"Successfully uploaded to Node server. Persistent URL: {persistent_url}")
+    except requests.exceptions.RequestException as upload_err:
+        print(f"FAILED to upload GLB to Node server: {upload_err}")
+        if hasattr(upload_err, 'response') and upload_err.response is not None:
+            print(f"Node server response status: {upload_err.response.status_code}")
+            print(f"Node server response text: {upload_err.response.text}")
+        raise gr.Error(f"Failed to upload result to backend server: {upload_err}")
+    except Exception as e:
+        print(f"UNEXPECTED error during upload: {e}", exc_info=True)
+        raise gr.Error(f"Unexpected error during upload: {e}")
+    finally:
+        if os.path.exists(temp_glb_path):
+            print(f"Cleaning up temporary GLB: {temp_glb_path}")
+            os.remove(temp_glb_path)
+
+    if not persistent_url:
+        print("Failed to obtain a persistent URL for the generated model.")
+        raise gr.Error("Failed to obtain a persistent URL for the generated model.")
+
+    print(f"Returning persistent URL: {persistent_url}")
+    return persistent_url
+
 # Interfaz Gradio
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
@@ -206,6 +358,38 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[download_glb],
     )
 
+    # --- Add this section to explicitly register the API function for image to 3D ---
+    # These State components are placeholders for API-only inputs
+    api_image_inputs_state = gr.State(value=[]) # For List[str] of image_inputs
+    api_input_type_state = gr.State(value="url") # For input_type: "url", "filepath", or "base64"
+    api_model_description_state = gr.State(value="ImagenModel") # For model_description
+
+    with gr.Row(visible=False): # Hide this row in the UI
+        api_image_gen_trigger_btn = gr.Button("API Image-to-3D Trigger")
+    
+    # Output for the API call (can be a dummy Textbox)
+    api_image_gen_output_url = gr.Textbox(label="Generated Model URL (API)", visible=False)
+
+    api_image_gen_trigger_btn.click(
+         generate_model_from_images_and_upload,
+         inputs=[ # Order must match the Python function's parameters
+             api_image_inputs_state,
+             api_input_type_state,
+             seed, # UI component
+             ss_guidance_strength, # UI component
+             ss_sampling_steps, # UI component
+             slat_guidance_strength, # UI component
+             slat_sampling_steps, # UI component
+             multiimage_algo, # UI component
+             mesh_simplify, # UI component
+             texture_size, # UI component
+             api_model_description_state,
+         ],
+         outputs=[api_image_gen_output_url], 
+         api_name="generate_model_from_images_and_upload" # Critical: Register the API name
+    )
+    # --- End API registration section ---
+
 # Lanzar la aplicación Gradio
 if __name__ == "__main__":
     pipeline = TrellisImageTo3DPipeline.from_pretrained("cavargas10/TRELLIS")