app.py

"""
Infinigen Agents - AI-powered procedural 3D generation
Full version with Infinigen + Blender in Docker container
"""
import os
import sys
import gradio as gr
from typing import Dict, Any
from pathlib import Path

# Add infinigen to path (in Docker container)
sys.path.insert(0, "/app")
sys.path.insert(0, "/app/infinigen")

# HuggingFace token from Space secrets
HF_TOKEN = os.environ.get("HF_TOKEN")

# AI Model Configuration
AI_MODEL = os.environ.get("AI_MODEL", "huggingface")
HF_MODEL_ID = os.environ.get("HF_MODEL_ID", "openai/gpt-oss-20b")
HF_PROVIDER = os.environ.get("HF_PROVIDER", None)


def get_model():
    """Get configured AI model for agents."""
    if AI_MODEL == "huggingface":
        from pydantic_ai.models.huggingface import HuggingFaceModel
        from pydantic_ai.providers.huggingface import HuggingFaceProvider

        provider_kwargs = {"api_key": HF_TOKEN}
        if HF_PROVIDER:
            provider_kwargs["provider_name"] = HF_PROVIDER

        return HuggingFaceModel(HF_MODEL_ID, provider=HuggingFaceProvider(**provider_kwargs))
    else:
        return f"openai:gpt-4o-mini"


def compose_scene(scene_type: str, seed: int, complexity: str) -> Dict[str, Any]:
    """Compose a scene using AI agent."""
    try:
        from pydantic_ai import Agent

        agent = Agent(
            get_model(),
            system_prompt=f"""You are a scene composer for Infinigen.
            Create a {complexity} complexity {scene_type} scene with seed {seed}.
            Respond with JSON containing: scene_type, seed, assets, lighting, camera."""
        )

        result = agent.run_sync(f"Create a {scene_type} scene")
        return {
            "success": True,
            "scene_type": scene_type,
            "seed": seed,
            "complexity": complexity,
            "result": str(result.data)
        }
    except Exception as e:
        return {"success": False, "error": str(e)}


def generate_terrain(terrain_type: str, seed: int, resolution: int) -> Dict[str, Any]:
    """Generate terrain using AI agent."""
    try:
        from pydantic_ai import Agent

        agent = Agent(
            get_model(),
            system_prompt=f"""You are a terrain engineer for Infinigen.
            Generate {terrain_type} terrain with resolution {resolution}.
            Respond with terrain parameters: heightmap settings, erosion, materials."""
        )

        result = agent.run_sync(f"Generate {terrain_type} terrain")
        return {
            "success": True,
            "terrain_type": terrain_type,
            "seed": seed,
            "resolution": resolution,
            "result": str(result.data)
        }
    except Exception as e:
        return {"success": False, "error": str(e)}


def get_recommendations(scene_type: str) -> str:
    """Get AI recommendations for scene generation."""
    try:
        from pydantic_ai import Agent

        agent = Agent(
            get_model(),
            system_prompt="""You are an expert on Infinigen procedural generation.
            Provide recommendations for assets, terrain, lighting, and camera setup."""
        )

        result = agent.run_sync(f"Recommend settings for a {scene_type} scene in Infinigen")
        return str(result.data)
    except Exception as e:
        return f"Error: {e}"


# Gradio Interface
with gr.Blocks(title="Infinigen Agents") as demo:
    gr.Markdown("""
    # 🌍 Infinigen Agents
    **AI-powered procedural 3D world generation**

    Full version with Infinigen + Blender - Using HuggingFace Inference API
    """)

    with gr.Tab("Scene Composer"):
        with gr.Row():
            scene_type = gr.Dropdown(
                ["forest", "desert", "mountain", "canyon", "coast", "kitchen", "living_room"],
                label="Scene Type",
                value="forest"
            )
            scene_seed = gr.Number(label="Seed", value=42)
            complexity = gr.Dropdown(["low", "medium", "high"], label="Complexity", value="medium")

        compose_btn = gr.Button("🎬 Compose Scene", variant="primary")
        scene_output = gr.JSON(label="Scene Result")

        compose_btn.click(compose_scene, [scene_type, scene_seed, complexity], scene_output)

    with gr.Tab("Terrain Engineer"):
        with gr.Row():
            terrain_type = gr.Dropdown(
                ["mountain", "canyon", "cliff", "mesa", "river", "volcano", "coast", "plain"],
                label="Terrain Type",
                value="mountain"
            )
            terrain_seed = gr.Number(label="Seed", value=42)
            resolution = gr.Slider(128, 2048, value=512, step=128, label="Resolution")

        terrain_btn = gr.Button("🏔️ Generate Terrain", variant="primary")
        terrain_output = gr.JSON(label="Terrain Result")

        terrain_btn.click(generate_terrain, [terrain_type, terrain_seed, resolution], terrain_output)

    with gr.Tab("AI Recommendations"):
        rec_scene_type = gr.Dropdown(
            ["forest", "desert", "mountain", "canyon", "coast"],
            label="Scene Type",
            value="forest"
        )
        rec_btn = gr.Button("💡 Get Recommendations", variant="primary")
        rec_output = gr.Textbox(label="AI Recommendations", lines=10)

        rec_btn.click(get_recommendations, rec_scene_type, rec_output)

    gr.Markdown(f"""
    ---
    ### Configuration
    - **AI Model**: {AI_MODEL}
    - **HF Model**: {HF_MODEL_ID}
    - **Provider**: {HF_PROVIDER or 'auto'}

    ### MCP Server
    ```json
    {{
      "mcpServers": {{
        "infinigen-agents": {{
          "url": "https://dev-bjoern-infinigen-agents.hf.space/gradio_api/mcp/sse"
        }}
      }}
    }}
    ```
    """)


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, mcp_server=True)