302 lines
9.1 KiB
Go
302 lines
9.1 KiB
Go
package tasks
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"jiggablend/internal/runner/blender"
|
|
"jiggablend/internal/runner/workspace"
|
|
"jiggablend/pkg/scripts"
|
|
"jiggablend/pkg/types"
|
|
)
|
|
|
|
// RenderProcessor handles render tasks.
|
|
type RenderProcessor struct{}
|
|
|
|
// NewRenderProcessor creates a new render processor.
|
|
func NewRenderProcessor() *RenderProcessor {
|
|
return &RenderProcessor{}
|
|
}
|
|
|
|
// Process executes a render task.
|
|
func (p *RenderProcessor) Process(ctx *Context) error {
|
|
ctx.Info(fmt.Sprintf("Starting task: job %d, frame %d, format: %s",
|
|
ctx.JobID, ctx.Frame, ctx.GetOutputFormat()))
|
|
log.Printf("Processing task %d: job %d, frame %d", ctx.TaskID, ctx.JobID, ctx.Frame)
|
|
|
|
// Find .blend file
|
|
blendFile, err := workspace.FindFirstBlendFile(ctx.WorkDir)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to find blend file: %w", err)
|
|
}
|
|
|
|
// Get Blender binary
|
|
blenderBinary := "blender"
|
|
if version := ctx.GetBlenderVersion(); version != "" {
|
|
ctx.Info(fmt.Sprintf("Job requires Blender %s", version))
|
|
binaryPath, err := ctx.Blender.GetBinaryPath(version)
|
|
if err != nil {
|
|
ctx.Warn(fmt.Sprintf("Could not get Blender %s, using system blender: %v", version, err))
|
|
} else {
|
|
blenderBinary = binaryPath
|
|
ctx.Info(fmt.Sprintf("Using Blender binary: %s", blenderBinary))
|
|
}
|
|
} else {
|
|
ctx.Info("No Blender version specified, using system blender")
|
|
}
|
|
|
|
// Create output directory
|
|
outputDir := filepath.Join(ctx.WorkDir, "output")
|
|
if err := os.MkdirAll(outputDir, 0755); err != nil {
|
|
return fmt.Errorf("failed to create output directory: %w", err)
|
|
}
|
|
|
|
// Create home directory for Blender inside workspace
|
|
blenderHome := filepath.Join(ctx.WorkDir, "home")
|
|
if err := os.MkdirAll(blenderHome, 0755); err != nil {
|
|
return fmt.Errorf("failed to create Blender home directory: %w", err)
|
|
}
|
|
|
|
// Determine render format
|
|
outputFormat := ctx.GetOutputFormat()
|
|
renderFormat := outputFormat
|
|
if outputFormat == "EXR_264_MP4" || outputFormat == "EXR_AV1_MP4" || outputFormat == "EXR_VP9_WEBM" {
|
|
renderFormat = "EXR" // Use EXR for maximum quality
|
|
}
|
|
|
|
// Create render script
|
|
if err := p.createRenderScript(ctx, renderFormat); err != nil {
|
|
return err
|
|
}
|
|
|
|
// Render
|
|
ctx.Info(fmt.Sprintf("Starting Blender render for frame %d...", ctx.Frame))
|
|
if err := p.runBlender(ctx, blenderBinary, blendFile, outputDir, renderFormat, blenderHome); err != nil {
|
|
ctx.Error(fmt.Sprintf("Blender render failed: %v", err))
|
|
return err
|
|
}
|
|
|
|
// Verify output
|
|
if _, err := p.findOutputFile(ctx, outputDir, renderFormat); err != nil {
|
|
ctx.Error(fmt.Sprintf("Output verification failed: %v", err))
|
|
return err
|
|
}
|
|
ctx.Info(fmt.Sprintf("Blender render completed for frame %d", ctx.Frame))
|
|
|
|
return nil
|
|
}
|
|
|
|
func (p *RenderProcessor) createRenderScript(ctx *Context, renderFormat string) error {
|
|
formatFilePath := filepath.Join(ctx.WorkDir, "output_format.txt")
|
|
renderSettingsFilePath := filepath.Join(ctx.WorkDir, "render_settings.json")
|
|
|
|
// Build unhide code conditionally
|
|
unhideCode := ""
|
|
if ctx.ShouldUnhideObjects() {
|
|
unhideCode = scripts.UnhideObjects
|
|
}
|
|
|
|
// Load template and replace placeholders
|
|
scriptContent := scripts.RenderBlenderTemplate
|
|
scriptContent = strings.ReplaceAll(scriptContent, "{{UNHIDE_CODE}}", unhideCode)
|
|
scriptContent = strings.ReplaceAll(scriptContent, "{{FORMAT_FILE_PATH}}", fmt.Sprintf("%q", formatFilePath))
|
|
scriptContent = strings.ReplaceAll(scriptContent, "{{RENDER_SETTINGS_FILE}}", fmt.Sprintf("%q", renderSettingsFilePath))
|
|
|
|
scriptPath := filepath.Join(ctx.WorkDir, "enable_gpu.py")
|
|
if err := os.WriteFile(scriptPath, []byte(scriptContent), 0644); err != nil {
|
|
errMsg := fmt.Sprintf("failed to create GPU enable script: %v", err)
|
|
ctx.Error(errMsg)
|
|
return errors.New(errMsg)
|
|
}
|
|
|
|
// Write output format
|
|
outputFormat := ctx.GetOutputFormat()
|
|
ctx.Info(fmt.Sprintf("Writing output format '%s' to format file", outputFormat))
|
|
if err := os.WriteFile(formatFilePath, []byte(outputFormat), 0644); err != nil {
|
|
errMsg := fmt.Sprintf("failed to create format file: %v", err)
|
|
ctx.Error(errMsg)
|
|
return errors.New(errMsg)
|
|
}
|
|
|
|
// Write render settings if available
|
|
if ctx.Metadata != nil && ctx.Metadata.RenderSettings.EngineSettings != nil {
|
|
settingsJSON, err := json.Marshal(ctx.Metadata.RenderSettings)
|
|
if err == nil {
|
|
if err := os.WriteFile(renderSettingsFilePath, settingsJSON, 0644); err != nil {
|
|
ctx.Warn(fmt.Sprintf("Failed to write render settings file: %v", err))
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (p *RenderProcessor) runBlender(ctx *Context, blenderBinary, blendFile, outputDir, renderFormat, blenderHome string) error {
|
|
scriptPath := filepath.Join(ctx.WorkDir, "enable_gpu.py")
|
|
|
|
args := []string{"-b", blendFile, "--python", scriptPath}
|
|
if ctx.ShouldEnableExecution() {
|
|
args = append(args, "--enable-autoexec")
|
|
}
|
|
|
|
// Output pattern
|
|
outputPattern := filepath.Join(outputDir, fmt.Sprintf("frame_####.%s", strings.ToLower(renderFormat)))
|
|
outputAbsPattern, _ := filepath.Abs(outputPattern)
|
|
args = append(args, "-o", outputAbsPattern)
|
|
|
|
args = append(args, "-f", fmt.Sprintf("%d", ctx.Frame))
|
|
|
|
// Wrap with xvfb-run
|
|
xvfbArgs := []string{"-a", "-s", "-screen 0 800x600x24", blenderBinary}
|
|
xvfbArgs = append(xvfbArgs, args...)
|
|
cmd := exec.Command("xvfb-run", xvfbArgs...)
|
|
cmd.Dir = ctx.WorkDir
|
|
|
|
// Set up environment with custom HOME directory
|
|
env := os.Environ()
|
|
// Remove existing HOME if present and add our custom one
|
|
newEnv := make([]string, 0, len(env)+1)
|
|
for _, e := range env {
|
|
if !strings.HasPrefix(e, "HOME=") {
|
|
newEnv = append(newEnv, e)
|
|
}
|
|
}
|
|
newEnv = append(newEnv, fmt.Sprintf("HOME=%s", blenderHome))
|
|
cmd.Env = newEnv
|
|
|
|
// Set up pipes
|
|
stdoutPipe, err := cmd.StdoutPipe()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create stdout pipe: %w", err)
|
|
}
|
|
|
|
stderrPipe, err := cmd.StderrPipe()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create stderr pipe: %w", err)
|
|
}
|
|
|
|
if err := cmd.Start(); err != nil {
|
|
return fmt.Errorf("failed to start blender: %w", err)
|
|
}
|
|
|
|
// Track process
|
|
ctx.Processes.Track(ctx.TaskID, cmd)
|
|
defer ctx.Processes.Untrack(ctx.TaskID)
|
|
|
|
// Stream stdout
|
|
stdoutDone := make(chan bool)
|
|
go func() {
|
|
defer close(stdoutDone)
|
|
scanner := bufio.NewScanner(stdoutPipe)
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
if line != "" {
|
|
shouldFilter, logLevel := blender.FilterLog(line)
|
|
if !shouldFilter {
|
|
ctx.Log(logLevel, line)
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
|
|
// Stream stderr
|
|
stderrDone := make(chan bool)
|
|
go func() {
|
|
defer close(stderrDone)
|
|
scanner := bufio.NewScanner(stderrPipe)
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
if line != "" {
|
|
shouldFilter, logLevel := blender.FilterLog(line)
|
|
if !shouldFilter {
|
|
if logLevel == types.LogLevelInfo {
|
|
logLevel = types.LogLevelWarn
|
|
}
|
|
ctx.Log(logLevel, line)
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
|
|
// Wait for completion
|
|
err = cmd.Wait()
|
|
<-stdoutDone
|
|
<-stderrDone
|
|
|
|
if err != nil {
|
|
if exitErr, ok := err.(*exec.ExitError); ok {
|
|
if exitErr.ExitCode() == 137 {
|
|
return errors.New("Blender was killed due to excessive memory usage (OOM)")
|
|
}
|
|
}
|
|
return fmt.Errorf("blender failed: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (p *RenderProcessor) findOutputFile(ctx *Context, outputDir, renderFormat string) (string, error) {
|
|
entries, err := os.ReadDir(outputDir)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to read output directory: %w", err)
|
|
}
|
|
|
|
ctx.Info("Checking output directory for files...")
|
|
|
|
// Try exact match first
|
|
expectedFile := filepath.Join(outputDir, fmt.Sprintf("frame_%04d.%s", ctx.Frame, strings.ToLower(renderFormat)))
|
|
if _, err := os.Stat(expectedFile); err == nil {
|
|
ctx.Info(fmt.Sprintf("Found output file: %s", filepath.Base(expectedFile)))
|
|
return expectedFile, nil
|
|
}
|
|
|
|
// Try without zero padding
|
|
altFile := filepath.Join(outputDir, fmt.Sprintf("frame_%d.%s", ctx.Frame, strings.ToLower(renderFormat)))
|
|
if _, err := os.Stat(altFile); err == nil {
|
|
ctx.Info(fmt.Sprintf("Found output file: %s", filepath.Base(altFile)))
|
|
return altFile, nil
|
|
}
|
|
|
|
// Try just frame number
|
|
altFile2 := filepath.Join(outputDir, fmt.Sprintf("%04d.%s", ctx.Frame, strings.ToLower(renderFormat)))
|
|
if _, err := os.Stat(altFile2); err == nil {
|
|
ctx.Info(fmt.Sprintf("Found output file: %s", filepath.Base(altFile2)))
|
|
return altFile2, nil
|
|
}
|
|
|
|
// Search through all files
|
|
for _, entry := range entries {
|
|
if !entry.IsDir() {
|
|
fileName := entry.Name()
|
|
if strings.Contains(fileName, "%04d") || strings.Contains(fileName, "%d") {
|
|
ctx.Warn(fmt.Sprintf("Skipping file with literal pattern: %s", fileName))
|
|
continue
|
|
}
|
|
frameStr := fmt.Sprintf("%d", ctx.Frame)
|
|
frameStrPadded := fmt.Sprintf("%04d", ctx.Frame)
|
|
if strings.Contains(fileName, frameStrPadded) ||
|
|
(strings.Contains(fileName, frameStr) && strings.HasSuffix(strings.ToLower(fileName), strings.ToLower(renderFormat))) {
|
|
outputFile := filepath.Join(outputDir, fileName)
|
|
ctx.Info(fmt.Sprintf("Found output file: %s", fileName))
|
|
return outputFile, nil
|
|
}
|
|
}
|
|
}
|
|
|
|
// Not found
|
|
fileList := []string{}
|
|
for _, entry := range entries {
|
|
if !entry.IsDir() {
|
|
fileList = append(fileList, entry.Name())
|
|
}
|
|
}
|
|
return "", fmt.Errorf("output file not found: %s\nFiles in output directory: %v", expectedFile, fileList)
|
|
}
|