Hang and crash when building computational graph

[jerr060599]

Hi,

Diffvg appears to have issues when multiple instances of RenderFunction.apply are to be added into the computational graph. If a second instance is called while an instance already exists under the computational graph, it will first hang and then eventually crash.

It does not output any errors to the console. It only hangs before crashing.

Here a simple script that demonstrates the behavior. It is a simple polygon matching example, except with two keyframes. It tries to match the two keyframes to a slowly rotating square.

import pydiffvg
import torch
import skimage
import numpy as np
import math

# Use GPU if available
pydiffvg.set_use_gpu(torch.cuda.is_available())

canvas_width, canvas_height = 256, 256

# Initialize the some target frames
targets = [];
render = pydiffvg.RenderFunction.apply

for i in range(8):
    points = torch.tensor([[128 + 60 * math.cos(0.1 * i + 0.0 * math.pi), 128 + 60 * math.sin(0.1 * i + 0.0 * math.pi)],
                           [128 + 60 * math.cos(0.1 * i + 0.5 * math.pi), 128 + 60 * math.sin(0.1 * i + 0.5 * math.pi)],
                           [128 + 60 * math.cos(0.1 * i + 1.0 * math.pi), 128 + 60 * math.sin(0.1 * i + 1.0 * math.pi)],
                           [128 + 60 * math.cos(0.1 * i + 1.5 * math.pi), 128 + 60 * math.sin(0.1 * i + 1.5 * math.pi)]])
    polygon = pydiffvg.Polygon(points = points, is_closed = True)
    shapes = [polygon]
    polygon_group = pydiffvg.ShapeGroup(shape_ids = torch.tensor([0]),
                                    fill_color = torch.tensor([1.0, 1.0, 1.0, 1.0]))
    shape_groups = [polygon_group]
    scene_args = pydiffvg.RenderFunction.serialize_scene(\
                    canvas_width, canvas_height, shapes, shape_groups, 
                    output_type = pydiffvg.OutputType.sdf)
    img = render(256, # width
             256, # height
             2,   # num_samples_x
             2,   # num_samples_y
             0,   # seed
             None, # background_image
             *scene_args)
    img = img / 256
    pydiffvg.imwrite(img.cpu(), 'results/test2/target_{}.png'.format(i), gamma=2.2)
    targets.append(img.clone())

# Setup the scene
# Normalize points for easier learning rate
keyframes = []
keyframes.append(torch.tensor([[(128 + 20) / 256.0, (128 - 20) / 256.0],
                               [(128 + 20) / 256.0, (128 + 20) / 256.0],
                               [(128 - 20) / 256.0, (128 + 20) / 256.0],
                               [(128 - 20) / 256.0, (128 - 20) / 256.0]],
                              requires_grad = True))
keyframes.append(torch.tensor([[(128 + 30) / 256.0, (128 - 30) / 256.0],
                               [(128 + 30) / 256.0, (128 + 30) / 256.0],
                               [(128 - 30) / 256.0, (128 + 30) / 256.0],
                               [(128 - 30) / 256.0, (128 - 30) / 256.0]],
                              requires_grad = True))

polygon.points = keyframes[0] * 256
scene_args = pydiffvg.RenderFunction.serialize_scene(\
                canvas_width, canvas_height, shapes, shape_groups,
                output_type = pydiffvg.OutputType.sdf)
img = render(256, # width
             256, # height
             2,   # num_samples_x
             2,   # num_samples_y
             1,   # seed
             None, # background_image
             *scene_args)
img = img / 256
pydiffvg.imwrite(img.cpu(), 'results/test2/init.png', gamma=2.2)

# Optimizer. This is so nice. Much cleaner than in c++
optimizer = torch.optim.Adam(keyframes, lr=1e-2)

# Iterate and optimize
for t in range(100):
    print('iteration:', t)
    # Reset iteration
    optimizer.zero_grad()

    loss = torch.tensor(0.0, requires_grad = True)

    # Render current scene
    polygon.points = keyframes[0] * 256
    scene_args = pydiffvg.RenderFunction.serialize_scene(\
                    canvas_width, canvas_height, shapes, shape_groups,
                    output_type = pydiffvg.OutputType.sdf)
    img = render(256,   # width
                 256,   # height
                 2,  # num_samples_x
                 2,  # num_samples_y
                 t+1,   # seed
                 None,
                 *scene_args)
    img = img / 256
    pydiffvg.imwrite(img.cpu(), 'results/test2/iter_{}.png'.format(t), gamma=2.2)

    # Compute loss
    for i, tar in enumerate(targets):
        k = i / (len(targets) - 1)
        print('loss_{}'.format(k))

        polygon.points = ((1 - k) * keyframes[0] + k * keyframes[1]) * 256
        scene_args = pydiffvg.RenderFunction.serialize_scene(\
                        canvas_width, canvas_height, shapes, shape_groups,
                        output_type = pydiffvg.OutputType.sdf)
        print('scene_{}'.format(k))

        img = render(256,   # width
                     256,   # height
                     2,  # num_samples_x
                     2,  # num_samples_y
                     t+1,   # seed
                     None,
                     *scene_args)
        print('render_{}'.format(k))

        img = img / 256
        loss = loss + (img - tar).pow(2).sum()
        print('fin_{}'.format(k))

    print('loss:', loss.item())
    loss.backward()

     # Take a gradient descent step.
    optimizer.step()

# Render the final result.
for i, tar in enumerate(targets):
    k = i / (len(targets) - 1)

    polygon.points = ((1 - k) * keyframes[0] + k * keyframes[1]) * 256
    scene_args = pydiffvg.RenderFunction.serialize_scene(\
                    canvas_width, canvas_height, shapes, shape_groups,
                    output_type = pydiffvg.OutputType.sdf)
    img = render(256,   # width
                 256,   # height
                 2,  # num_samples_x
                 2,  # num_samples_y
                 i+1,   # seed
                 None,
                 *scene_args)
    img = img / 256
    pydiffvg.imwrite(img.cpu(), 'results/test2/final_{}.png'.format(i), gamma=2.2)

from subprocess import call
call(["ffmpeg", "-framerate", "24", "-i",
    "results/test2/iter_%d.png", "-vb", "20M",
    "results/test2/out.mp4"])

I am using Windows 11 and a NVidia GPU.