app.py

#!/usr/bin/env python3

import gradio as gr
import subprocess
from pathlib import Path
import torch
import os
import logging
import sys
import spaces

# Logging konfigurieren
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

class InfinigenManager:
    def __init__(self, base_dir="infinigen", output_dir="outputs", blender_version="4.3.2"):
        """Initialisiert den InfinigenManager mit Basisverzeichnis, Ausgabeverzeichnis und Blender-Version."""
        self.base_dir = Path(base_dir)
        self.output_dir = Path(output_dir)
        self.blender_version = blender_version
        self.blender_bin = None
        self.blender_python = None
        self.initialized = False
        logger.info("Initialisiere InfinigenManager...")
        self._initialize()

    def _run_command(self, command, error_msg="Fehler bei Befehlsausführung", timeout=None):
        """Führt Befehle einheitlich aus und behandelt Fehler."""
        try:
            result = subprocess.run(
                command,
                shell=isinstance(command, str),
                capture_output=True,
                text=True,
                check=True,
                timeout=timeout
            )
            logger.debug(f"STDOUT: {result.stdout}")
            return result
        except subprocess.CalledProcessError as e:
            logger.error(f"{error_msg}: {e.stderr}")
            raise
        except subprocess.TimeoutExpired as e:
            logger.error(f"{error_msg} (Timeout): {e.stderr}")
            raise

    def _initialize(self):
        """Führt die zentrale Initialisierung durch."""
        try:
            self._setup_blender()
            self._setup_infinigen()
            self._check_cuda()
            self.initialized = True
            logger.info("Initialisierung erfolgreich abgeschlossen.")
        except Exception as e:
            logger.error(f"Initialisierung fehlgeschlagen: {e}")
            self.initialized = False

    def _setup_blender(self):
        """Richtet Blender und dessen Python-Interpreter ein."""
        logger.info("Starte Blender-Setup...")
        self.blender_bin = self._find_blender() or self._install_latest_blender()
        self.blender_python = self._find_blender_python_with_grep()
        if not self.blender_python:
            raise RuntimeError("Blender Python-Interpreter nicht gefunden!")
        logger.info(f"Blender Python gefunden: {self.blender_python}")

    def _find_blender(self):
        """Sucht nach einer vorhandenen Blender-Installation."""
        blender_paths = [
            f"/home/user/blender-{self.blender_version}-linux-x64/blender",
            "/usr/bin/blender"
        ]
        for bin_path in blender_paths:
            if os.path.exists(bin_path):
                try:
                    version_result = self._run_command([bin_path, "-v"], "Kann Blender-Version nicht prüfen", timeout=10)
                    version = next((line.split()[1] for line in version_result.stdout.split("\n")
                                   if "Blender" in line and "." in line), None)
                    logger.info(f"Vorhandene Blender-Version: {version or 'unbekannt'}")
                    if version and version >= self.blender_version:
                        return bin_path
                except Exception as e:
                    logger.warning(f"Fehler bei Blender-Prüfung ({bin_path}): {e}")
        return None

    def _install_latest_blender(self):
        """Installiert Blender 4.3.2, falls nicht vorhanden."""
        logger.info("Installiere Blender 4.3.2...")
        url = "https://download.blender.org/release/Blender4.3/blender-4.3.2-linux-x64.tar.xz"
        tar_path = "blender-4.3.2-linux-x64.tar.xz"
        blender_dir = Path(f"/home/user/blender-{self.blender_version}-linux-x64")
        
        self._run_command(f"wget -q {url}", "Fehler beim Herunterladen von Blender", timeout=60)
        self._run_command(f"tar -xvf {tar_path} -C /home/user", "Fehler beim Entpacken von Blender", timeout=60)
        os.remove(tar_path)
        return blender_dir / "blender"

    def _find_blender_python_with_grep(self):
        """Sucht den Blender-Python-Pfad mit find."""
        base_path = Path(self.blender_bin).parent.parent
        grep_cmd = f"find {base_path} -type f -path '*/python/bin/python3*' -executable"
        try:
            result = self._run_command(grep_cmd, "Fehler bei der Python-Suche mit find")
            python_paths = result.stdout.strip().split("\n")
            for path in python_paths:
                if path:
                    version_result = self._run_command([path, "--version"], "Fehler bei Python-Version", timeout=5)
                    logger.info(f"Python-Version: {version_result.stdout.strip()}")
                    return path
            logger.warning("Kein passender Python-Pfad in find-Ausgabe gefunden.")
            return None
        except Exception as e:
            logger.error(f"Fehler bei der Suche nach Blender-Python: {e}")
            return None

    def _check_cuda(self):
        """Prüft die CUDA-Verfügbarkeit für GPU-Nutzung."""
        cuda_available = torch.cuda.is_available()
        logger.info(f"CUDA {'verfügbar' if cuda_available else 'nicht verfügbar'}: {torch.cuda.get_device_name(0) if cuda_available else 'N/A'}")

    def _setup_infinigen(self):
        """Richtet Infinigen ein, falls nicht installiert."""
        if not self.blender_python:
            logger.error("Kann Infinigen nicht installieren: Blender Python fehlt.")
            return

        logger.info("Starte Infinigen-Setup...")
        if not self.base_dir.exists():
            self._run_command(f"git clone https://github.com/princeton-vl/infinigen.git {self.base_dir}",
                            "Fehler beim Klonen von Infinigen", timeout=30)

        # Installiere Infinigen mit Abhängigkeiten
        self._run_command([
            self.blender_python, "-m", "pip", "install", "-e", f"{self.base_dir}[terrain,vis]", "--user"
        ], "Fehler bei der Installation von Infinigen", timeout=60)

        # Überprüfe, ob das datagen-Modul verfügbar ist
        if not self._is_datagen_available():
            raise RuntimeError("Das 'datagen'-Modul von Infinigen ist nicht verfügbar. Überprüfe die Installation.")

    def _is_datagen_available(self):
        """Prüft, ob das 'datagen'-Modul von Infinigen verfügbar ist."""
        try:
            self._run_command([self.blender_python, "-c", "import infinigen.datagen"], timeout=10)
            return True
        except Exception as e:
            logger.error(f"Fehler beim Import von 'infinigen.datagen': {e}")
            return False

    @spaces.GPU
    def generate_scene(self, seed, configs=None, pipeline_configs=None):
        """Generiert eine Szene mit Infinigen auf der GPU."""
        if not self.initialized or not self.blender_python:
            return "Fehler: Infinigen nicht initialisiert!"

        logger.info(f"Generiere Szene mit Seed: {seed}")
        self.output_dir.mkdir(exist_ok=True)
        configs = configs or ["infinigen_examples/configs/desert.gin", "infinigen_examples/configs/simple.gin"]
        pipeline_configs = pipeline_configs or [
            "infinigen_examples/configs/local_16GB.gin",
            "infinigen_examples/configs/monocular.gin",
            "infinigen_examples/configs/blender_gt.gin"
        ]

        command = [
            self.blender_python, "-m", "infinigen.datagen.manage_jobs",
            "--output_folder", str(self.output_dir),
            "--num_scenes", "1",
            "--specific_seed", str(int(seed)),
            "--configs"
        ] + configs + ["--pipeline_configs"] + pipeline_configs

        try:
            self._run_command(command, "Fehler bei Szenengenerierung", timeout=300)
            output_path = self.output_dir / "0000000000.png"
            return str(output_path) if output_path.exists() else "Fehler: Bild nicht gefunden!"
        except Exception as e:
            return f"Fehler bei Szenengenerierung: {str(e)}"

# Manager initialisieren
logger.info("Starte Manager-Initialisierung...")
manager = InfinigenManager()

# Gradio-Oberfläche definieren
with gr.Blocks(title="Infinigen Demo") as demo:
    gr.Markdown("## Infinigen Scene Generator")
    seed_input = gr.Number(label="Seed", value=0, precision=0)
    output_image = gr.Image(label="Generierte Szene")
    generate_button = gr.Button("Szene generieren")
    
    generate_button.click(fn=manager.generate_scene, inputs=[seed_input], outputs=[output_image])

logger.info("Starte Gradio-Oberfläche...")
demo.launch()