initial demo release

app.py

@@ -1,7 +1,263 @@
 import gradio as gr
+import os
+import sys
+import subprocess
+from huggingface_hub import snapshot_download, HfFolder
+import random # Import random for seed generation
 
-def greet(name):
-    return "Hello " + name + "!!"
+# --- Repo Setup --- 
+DEFAULT_REPO_DIR = "./TripoSG-repo" # Directory to clone into if not using local path
+REPO_GIT_URL = "github.com/VAST-AI-Research/TripoSG.git" # Base URL without schema/token
+BRANCH = "scribble"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+code_source_path = None
+
+# Option 1: Use local path if TRIPOSG_CODE_PATH env var is set
+local_code_path = os.environ.get("TRIPOSG_CODE_PATH")
+if local_code_path:
+    print(f"Attempting to use local code path specified by TRIPOSG_CODE_PATH: {local_code_path}")
+    # Basic check: does it exist and seem like a git repo (has .git)?
+    if os.path.isdir(local_code_path) and os.path.isdir(os.path.join(local_code_path, ".git")):
+        code_source_path = os.path.abspath(local_code_path)
+        print(f"Using local TripoSG code directory: {code_source_path}")
+        # You might want to add a check here to verify the branch is correct, e.g.:
+        # try:
+        #     current_branch = subprocess.run(["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=code_source_path, check=True, capture_output=True, text=True).stdout.strip()
+        #     if current_branch != BRANCH:
+        #         print(f"Warning: Local repo is on branch '{current_branch}', expected '{BRANCH}'. Attempting checkout...")
+        #         subprocess.run(["git", "checkout", BRANCH], cwd=code_source_path, check=True)
+        # except Exception as e:
+        #     print(f"Warning: Could not verify or checkout branch '{BRANCH}' in {code_source_path}: {e}")
+    else:
+        print(f"Warning: TRIPOSG_CODE_PATH '{local_code_path}' not found or not a valid git repository directory. Falling back to cloning.")
+
+# Option 2: Clone from GitHub (if local path not used or invalid)
+if not code_source_path:
+    repo_url_to_clone = f"https://{REPO_GIT_URL}"
+    github_token = os.environ.get("GITHUB_TOKEN")
+    if github_token:
+        print("Using GITHUB_TOKEN for repository cloning.")
+        repo_url_to_clone = f"https://{github_token}@{REPO_GIT_URL}"
+    else:
+        print("No GITHUB_TOKEN found. Using public HTTPS for cloning.")
+
+    repo_target_dir = os.path.abspath(DEFAULT_REPO_DIR)
+    if not os.path.exists(repo_target_dir):
+        print(f"Cloning TripoSG repository ({BRANCH} branch) into {repo_target_dir}...")
+        try:
+            subprocess.run(["git", "clone", "--branch", BRANCH, "--depth", "1", repo_url_to_clone, repo_target_dir], check=True)
+            code_source_path = repo_target_dir
+            print("Repository cloned successfully.")
+        except subprocess.CalledProcessError as e:
+            print(f"Error cloning repository: {e}")
+            print("Please ensure the URL is correct, the branch '{BRANCH}' exists, and you have access rights (or provide a GITHUB_TOKEN).")
+            sys.exit(1)
+        except Exception as e:
+            print(f"An unexpected error occurred during cloning: {e}")
+            sys.exit(1)
+    else:
+        print(f"Directory {repo_target_dir} already exists. Assuming it contains the correct code/branch.")
+        # Optional: Add checks here like git pull or verifying the branch
+        code_source_path = repo_target_dir
+
+if not code_source_path:
+    print("Error: Could not determine TripoSG code source path.")
+    sys.exit(1)
+
+# Add repo to Python path
+sys.path.insert(0, code_source_path) # Use the determined absolute path
+print(f"Added {code_source_path} to sys.path")
+# --- End Repo Setup ---
+
+# --- ZeroGPU Setup ---
+ENABLE_ZEROGPU = os.environ.get("ENABLE_ZEROGPU", "false").lower() in ("true", "1", "t")
+print(f"ZeroGPU Enabled: {ENABLE_ZEROGPU}")
+# --- End ZeroGPU Setup ---
+
+if ENABLE_ZEROGPU:
+    import spaces # Import spaces for ZeroGPU
+from PIL import Image
+import numpy as np
+import torch
+from triposg.pipelines.pipeline_triposg_scribble import TripoSGScribblePipeline
+import tempfile
+
+# --- Weight Loading Logic ---
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if HF_TOKEN:
+    HfFolder.save_token(HF_TOKEN)
+HUGGING_FACE_REPO_ID = "VAST-AI/TripoSG-scribble"
+DEFAULT_CACHE_PATH = "./pretrained_weights/TripoSG-scribble"
+
+# Option 1: Use local path if WEIGHTS_PATH env var is set
+local_weights_path = os.environ.get("WEIGHTS_PATH")
+model_load_path = None
+
+if local_weights_path:
+    print(f"Attempting to load weights from local path specified by WEIGHTS_PATH: {local_weights_path}")
+    if os.path.isdir(local_weights_path):
+        model_load_path = local_weights_path
+        print(f"Using local weights directory: {model_load_path}")
+    else:
+        print(f"Warning: WEIGHTS_PATH '{local_weights_path}' not found or not a directory. Falling back to Hugging Face download.")
+
+# Option 2: Download from Hugging Face (if local path not used or invalid)
+if not model_load_path:
+    hf_token = os.environ.get("HF_TOKEN")
+    print(f"Attempting to download weights from Hugging Face repo: {HUGGING_FACE_REPO_ID}")
+    if hf_token:
+        print("Using Hugging Face token for download.")
+        auth_token = hf_token
+    else:
+        print("No Hugging Face token found. Attempting public download.")
+        auth_token = None
+    try:
+        model_load_path = snapshot_download(
+            repo_id=HUGGING_FACE_REPO_ID,
+            local_dir=DEFAULT_CACHE_PATH,
+            local_dir_use_symlinks=False, # Recommended for Spaces
+            token=auth_token,
+            # revision="main" # Specify branch/commit if needed
+        )
+        print(f"Weights downloaded/cached to: {model_load_path}")
+    except Exception as e:
+        print(f"Error downloading weights from Hugging Face: {e}")
+        print("Please ensure the repository exists and is accessible, or provide a valid WEIGHTS_PATH.")
+        sys.exit(1) # Exit if weights cannot be loaded
+
+# Load the pipeline using the determined path
+print(f"Loading pipeline from: {model_load_path}")
+pipe = TripoSGScribblePipeline.from_pretrained(model_load_path)
+pipe.to(dtype=torch.float16, device="cuda")
+print("Pipeline loaded.")
+# --- End Weight Loading Logic ---
+
+# Create a white background image and a transparent layer for drawing
+canvas_width, canvas_height = 512, 512
+initial_background = Image.new("RGB", (canvas_width, canvas_height), color="white")
+initial_layer = Image.new("RGBA", (canvas_width, canvas_height), color=(0, 0, 0, 0)) # Transparent layer
+# Prepare the initial value dictionary for ImageEditor
+initial_value = {
+    "background": initial_background,
+    "layers": [initial_layer], # Add the transparent layer
+    "composite": None
+}
+
+# --- ZeroGPU Setup ---
+# ... existing ZeroGPU setup ...
+
+MAX_SEED = np.iinfo(np.int32).max
+
+def get_random_seed():
+    return random.randint(0, MAX_SEED)
+
+# Apply decorator conditionally
+@spaces.GPU(duration=120) if ENABLE_ZEROGPU else lambda func: func
+def generate_3d(scribble_image_dict, prompt, scribble_confidence, seed): # Added seed parameter back
+    print("Generating 3D model...")
+    # Extract the composite image from the ImageEditor dictionary
+    if scribble_image_dict is None or scribble_image_dict.get("composite") is None:
+        print("No scribble image provided.")
+        return None # Return None if no image is provided
+
+    # --- Seed Handling ---
+    current_seed = int(seed)
+    print(f"Using seed: {current_seed}")
+    # --- End Seed Handling ---
+
+    # Get the composite image which includes the drawing
+    # The composite might be RGBA if a layer was involved, ensure RGB for processing
+    image = Image.fromarray(scribble_image_dict["composite"]).convert("RGB") 
+
+    # Preprocess the image: invert colors (black on white -> white on black)
+    image_np = np.array(image)
+    processed_image_np = 255 - image_np
+    processed_image = Image.fromarray(processed_image_np)
+    print("Image preprocessed.")
+
+    # Define fixed parameters
+    attn_scale_text = 1.0 # As per the example run.py
+
+    # Set the generator with the provided seed
+    generator = torch.Generator(device='cuda').manual_seed(current_seed)
+
+    # Run the pipeline
+    print("Running pipeline...")
+    out = pipe(
+        processed_image,
+        prompt=prompt,
+        num_tokens=512,
+        guidance_scale=0,
+        num_inference_steps=16,
+        attention_kwargs={
+            "cross_attention_scale": attn_scale_text,
+            "cross_attention_2_scale": scribble_confidence
+        },
+        generator=generator,
+        use_flash_decoder=False,
+        dense_octree_depth=8,
+        hierarchical_octree_depth=8
+    )
+    print("Pipeline finished.")
+
+    # Save the output mesh to a temporary file
+    if out.meshes and len(out.meshes) > 0:
+        # Create a temporary file with .glb extension
+        with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmpfile:
+            output_path = tmpfile.name
+        out.meshes[0].export(output_path)
+        print(f"Mesh saved to temporary file: {output_path}")
+        return output_path
+    else:
+        print("Pipeline did not generate any meshes.")
+        return None
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Scribble + Text to 3D Model Generator (TripoSG)")
+    gr.Markdown("Draw a scribble (black on white canvas), enter a text prompt, adjust confidence, set a seed, and generate a 3D model.") # Updated guidance
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.ImageEditor(
+                label="Scribble Input (Draw Black on White)",
+                value=initial_value,
+                image_mode="RGB",
+                brush=gr.Brush(default_color="#000000", color_mode="fixed", default_size=5), # Fixed small brush size
+                interactive=True,
+                eraser=gr.Brush(default_color="#FFFFFF", color_mode="fixed", default_size=20) # Fixed small eraser size
+            )
+            prompt_input = gr.Textbox(label="Prompt", placeholder="e.g., a cute cat wearing a hat")
+            confidence_input = gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.05, label="Scribble Confidence (attn_scale_image)")
+            seed_input = gr.Number(label="Seed", value=0, precision=0) # Added Seed input back
+            with gr.Row():
+                submit_button = gr.Button("Generate 3D Model", variant="primary", scale=1)
+                lucky_button = gr.Button("I'm Feeling Lucky", scale=1)
+        with gr.Column(scale=1):
+            model_output = gr.Model3D(label="Generated 3D Model", interactive=False)
+
+    # Define the inputs for the main generation function
+    gen_inputs = [image_input, prompt_input, confidence_input, seed_input]
+
+    submit_button.click(
+        fn=generate_3d,
+        inputs=gen_inputs, # Include seed_input
+        outputs=model_output
+    )
+
+    # Define inputs for the lucky button (same as main button for the final call)
+    lucky_gen_inputs = [image_input, prompt_input, confidence_input, seed_input]
+
+    lucky_button.click(
+        fn=get_random_seed, # First, get a random seed
+        inputs=[],
+        outputs=[seed_input] # Update the seed input field
+    ).then(
+        fn=generate_3d, # Then, generate the model
+        inputs=lucky_gen_inputs, # Use the updated seed from the input field
+        outputs=model_output
+    )
+
+# Launch with queue enabled if using ZeroGPU
+print("Launching Gradio interface...")
+demo.launch(share=False, server_name="0.0.0.0")
+print("Gradio interface launched.")

requirements.txt

@@ -0,0 +1,16 @@
+diffusers
+transformers==4.49.0
+einops
+huggingface_hub
+opencv-python
+trimesh==4.5.3
+omegaconf
+scikit-image
+numpy
+peft
+scipy==1.11.4
+jaxtyping
+typeguard
+ninja
+gltflib
+https://huggingface.co/spaces/VAST-AI/TripoSG/resolve/main/diso-0.1.4-cp310-cp310-linux_x86_64.whl?download=true