app.py

import gradio as gr
import subprocess
import os
import spaces
import glob
from pathlib import Path
import torch

class InfinigenManager:
    def __init__(self, base_dir="infinigen", output_dir="outputs", blender_version="4.3.2"):
        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
        print("Initialisiere InfinigenManager...")
        self._setup_blender()
        self._setup_infinigen()
        self._check_cuda()

    def _setup_blender(self):
        """Installiert Blender 4.3.2, falls nötig, und findet den Python-Interpreter."""
        print("Starte Blender-Setup...")
        try:
            # Prüfe vorhandene Blender-Version
            blender_bins = glob.glob("/usr/bin/blender")
            if blender_bins:
                self.blender_bin = blender_bins[0]
                version_result = subprocess.run([self.blender_bin, "-v"], capture_output=True, text=True, timeout=10)
                current_version = next((line.split()[1] for line in version_result.stdout.split("\n") if "Blender" in line and "." in line), None)
                print(f"Vorhandene Blender-Version: {current_version or 'unbekannt'}")
                if current_version and current_version >= self.blender_version:
                    print("Vorhandene Version ist aktuell genug.")
                else:
                    print("Installiere neueste Blender-Version...")
                    self._install_latest_blender()
            else:
                print("Kein Blender vorhanden, installiere neueste Version...")
                self._install_latest_blender()

            # Finde Blender-Python
            base_path = f"/home/user/blender-{self.blender_version}-linux-x64" if self.blender_bin == f"/home/user/blender-{self.blender_version}-linux-x64/blender" else "/usr/share/blender"
            python_patterns = [
                f"{base_path}/python/bin/python*",
                "/usr/share/blender/*/python/bin/python*",
                "/usr/lib/blender/*/python/bin/python*"
            ]
            for pattern in python_patterns:
                print(f"Suche nach Python mit Muster: {pattern}")
                python_bins = glob.glob(pattern, recursive=True)
                if python_bins:
                    self.blender_python = python_bins[0]
                    print(f"Blender Python gefunden: {self.blender_python}")
                    self.initialized = True
                    return
            
            raise RuntimeError("Blender Python-Interpreter nicht gefunden!")
        except Exception as e:
            print(f"Fehler bei Blender-Setup: {e}")
            self.blender_python = None

    def _install_latest_blender(self):
        """Installiert Blender 4.3.2 dynamisch zur Laufzeit."""
        try:
            print("Lade Blender 4.3.2 herunter...")
            subprocess.run([
                "wget", "https://download.blender.org/release/Blender4.3/blender-4.3.2-linux-x64.tar.xz"
            ], check=True, timeout=60)
            print("Entpacke Blender 4.3.2...")
            subprocess.run([
                "tar", "-xvf", "blender-4.3.2-linux-x64.tar.xz", "-C", "/home/user"
            ], check=True, timeout=60)
            self.blender_bin = "/home/user/blender-4.3.2-linux-x64/blender"
            print(f"Blender 4.3.2 installiert unter: {self.blender_bin}")
            os.remove("blender-4.3.2-linux-x64.tar.xz")
        except Exception as e:
            print(f"Fehler bei der Installation von Blender 4.3.2: {e}")
            raise

    def _check_cuda(self):
        """Prüft CUDA-Verfügbarkeit."""
        if torch.cuda.is_available():
            print(f"CUDA verfügbar: {torch.cuda.get_device_name(0)}")
        else:
            print("CUDA nicht verfügbar!")

    def _setup_infinigen(self):
        """Installiert Infinigen, falls möglich."""
        print("Starte Infinigen-Setup...")
        if not self.blender_python:
            print("Kann Infinigen nicht installieren: Blender Python fehlt.")
            return
        
        if not self.base_dir.exists():
            print("Klone Infinigen...")
            subprocess.run(["git", "clone", "https://github.com/princeton-vl/infinigen.git"], check=True, timeout=30)
        
        if not self._is_infinigen_installed():
            print("Installiere Infinigen...")
            subprocess.run([
                self.blender_python, "-m", "pip", "install", "-e", str(self.base_dir) + "[terrain,vis]", "--no-deps", "--user"
            ], check=True, timeout=60)
            print("Infinigen installiert.")

    def _is_infinigen_installed(self):
        """Prüft, ob Infinigen bereits installiert ist."""
        if not self.blender_python:
            return False
        try:
            subprocess.run([self.blender_python, "-c", "import infinigen"], check=True, capture_output=True, timeout=10)
            return True
        except subprocess.CalledProcessError:
            return False

    @spaces.GPU
    def generate_scene(self, seed, configs=None, pipeline_configs=None):
        """Generiert eine Szene mit Infinigen."""
        if not self.initialized or not self.blender_python:
            return "Fehler: Infinigen nicht initialisiert (Blender Python fehlt)!"
        
        print(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:
            result = subprocess.run(command, capture_output=True, text=True, check=True, timeout=300)
            print(f"STDOUT: {result.stdout}")
            print(f"STDERR: {result.stderr}")
            output_path = self.output_dir / "0000000000.png"
            return str(output_path) if output_path.exists() else f"Fehler: Bild nicht gefunden. STDERR: {result.stderr}"
        except subprocess.TimeoutExpired as e:
            return f"Fehler: Timeout bei Szenengenerierung: {e.stderr}"
        except subprocess.CalledProcessError as e:
            return f"Fehler: {e.stderr}"

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

# Gradio-Oberfläche
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])

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