diff options
| -rw-r--r-- | .gitignore | 1 | ||||
| -rw-r--r-- | main.py | 36 | ||||
| -rw-r--r-- | models/autoencoder.py | 2 | ||||
| -rw-r--r-- | requirements.txt | 3 | ||||
| -rw-r--r-- | src/dataloader.py | 7 | ||||
| -rw-r--r-- | src/sample.py | 195 | ||||
| -rw-r--r-- | src/train_autoencoder.py | 83 | ||||
| -rw-r--r-- | src/train_diffusion.py | 35 | ||||
| -rw-r--r-- | src/utils.py | 20 |
9 files changed, 253 insertions, 129 deletions
@@ -8,6 +8,7 @@ plots/ checkpoints/ # MLflow artifacts +mlruns/ mlflow.db # Python bytecode @@ -27,9 +27,9 @@ def main(): formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__ ) - + subparsers = parser.add_subparsers(dest="command", help="Available commands") - + # Train diffusion model train_parser = subparsers.add_parser("train", help="Train diffusion model") train_parser.add_argument("--epochs", type=int, default=10, help="Number of epochs") @@ -38,7 +38,7 @@ def main(): train_parser.add_argument("--attention", action="store_true", help="Use self-attention") train_parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint") train_parser.add_argument("--name", type=str, default=None, help="Run name") - + # Train autoencoder ae_parser = subparsers.add_parser("train-ae", help="Train autoencoder") ae_parser.add_argument("--epochs", type=int, default=10, help="Number of epochs") @@ -46,14 +46,14 @@ def main(): ae_parser.add_argument("--batch-size", type=int, default=64, help="Batch size") ae_parser.add_argument("--latent-channels", type=int, default=1, help="Latent channels") ae_parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint") - + # Sample from model sample_parser = subparsers.add_parser("sample", help="Generate samples") sample_parser.add_argument("--checkpoint", type=str, default="checkpoints/diffusion_latest.pt", help="Path to model checkpoint") sample_parser.add_argument("--n-samples", type=int, default=10, help="Samples per class") sample_parser.add_argument("--attention", action="store_true", help="Use attention in model") - + # Visualize diffusion viz_parser = subparsers.add_parser("visualize", help="Visualize diffusion process") viz_parser.add_argument("--checkpoint", type=str, default="checkpoints/diffusion_latest.pt", @@ -62,16 +62,16 @@ def main(): viz_parser.add_argument("--forward", action="store_true", help="Visualize forward diffusion") viz_parser.add_argument("--backward", action="store_true", help="Visualize backward diffusion") viz_parser.add_argument("--all", action="store_true", help="Run all visualizations") - + args = parser.parse_args() - + if args.command is None: parser.print_help() return - + # Set multiprocessing start method torch.multiprocessing.set_start_method("spawn", force=True) - + if args.command == "train": from src.train_diffusion import train_diffusion train_diffusion( @@ -82,7 +82,7 @@ def main(): checkpoint_path=args.checkpoint, run_name=args.name, ) - + elif args.command == "train-ae": from src.train_autoencoder import train_autoencoder train_autoencoder( @@ -92,22 +92,22 @@ def main(): latent_channels=args.latent_channels, checkpoint_path=args.checkpoint, ) - + elif args.command == "sample": import os from src.config import DEVICE from src.sample import generate_grid from src.utils import load_checkpoint from models import UNetMNIST - + model = UNetMNIST(use_attention=args.attention).to(DEVICE) if os.path.exists(args.checkpoint): model = load_checkpoint(model, os.path.basename(args.checkpoint)) else: print(f"Warning: Checkpoint {args.checkpoint} not found.") - + generate_grid(model, n_per_class=args.n_samples) - + elif args.command == "visualize": import os from src.config import DEVICE @@ -118,17 +118,17 @@ def main(): ) from src.utils import load_checkpoint from models import UNetMNIST - + model = UNetMNIST(use_attention=args.attention).to(DEVICE) if os.path.exists(args.checkpoint): model = load_checkpoint(model, os.path.basename(args.checkpoint)) - + if args.forward or args.all: visualize_forward_diffusion() - + if args.backward or args.all: visualize_backward_diffusion(model) - + if args.all: generate_grid(model) diff --git a/models/autoencoder.py b/models/autoencoder.py index 353ec89..abbc3be 100644 --- a/models/autoencoder.py +++ b/models/autoencoder.py @@ -60,7 +60,7 @@ class Autoencoder(AEModule): nn.Conv2d(64, latent_channels, kernel_size=3, padding=1), nn.ReLU(), - # 7x7 -> 8x8 (no activation - latent space should be unconstrained) + # 7x7 -> 8x8 nn.Conv2d(latent_channels, latent_channels, kernel_size=2, padding=1), ) diff --git a/requirements.txt b/requirements.txt index 34490ec..75c0fba 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,5 +2,6 @@ torch torchvision mlflow scipy -matplotlib +plotly +kaleido rich diff --git a/src/dataloader.py b/src/dataloader.py index a08ae27..d1e7fcb 100644 --- a/src/dataloader.py +++ b/src/dataloader.py @@ -38,7 +38,7 @@ class DiffusionDataset(Dataset): autoencoder: Optional autoencoder for latent diffusion """ - def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module = None): + def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module | None = None): super().__init__() self.dataset = dataset self.autoencoder = autoencoder @@ -87,7 +87,7 @@ class DiffusionDatasetX0(Dataset): autoencoder: Optional autoencoder for latent diffusion """ - def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module = None): + def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module | None = None): super().__init__() self.dataset = dataset self.autoencoder = autoencoder @@ -143,6 +143,7 @@ class AutoencoderDataset(Dataset): def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: data = self.dataset[idx][0].to(DEVICE) + # data = data * 2 - 1 # Normalize to [-1, 1] return data, data @@ -151,7 +152,7 @@ def get_diffusion_dataloader( batch_size: int = 64, shuffle: bool = True, num_workers: int = 4, - autoencoder: torch.nn.Module = None, + autoencoder: torch.nn.Module | None = None, ) -> DataLoader: """ Create a DataLoader for diffusion training. diff --git a/src/sample.py b/src/sample.py index 5b438ed..b2c570b 100644 --- a/src/sample.py +++ b/src/sample.py @@ -4,17 +4,18 @@ Generate samples from a trained model and visualize results. """ from models import UNetMNIST -from src.utils import ensure_dirs, tensor_to_image, load_checkpoint +from src.utils import ensure_dirs, save_plot, tensor_to_image, load_checkpoint from src.dataloader import get_mnist_dataset from src.diffusion import p_xt_1_xt_x0_pred, forward_diffusion, q_xt_x0 from src.config import ( DEVICE, DIFFU_STEPS, NB_CHANNEL, IMG_SIZE, NB_LABEL, - CHECKPOINT_DIR, PLOTS_DIR + CHECKPOINT_DIR, ) import os import torch import numpy as np -import matplotlib.pyplot as plt +import plotly.graph_objects as go +from plotly.subplots import make_subplots from rich.progress import track import sys @@ -65,12 +66,14 @@ def generate_samples( return x -def visualize_forward_diffusion(save_path: str | None = None): +def visualize_forward_diffusion(save_path: str | None = None) -> go.Figure: """ Visualize the forward diffusion process on a real image. Args: save_path: Path to save the visualization + Returns: + Plotly figure object """ ensure_dirs() @@ -90,42 +93,73 @@ def visualize_forward_diffusion(save_path: str | None = None): n_plots = 10 timesteps = np.linspace(1, DIFFU_STEPS, n_plots, dtype=int) - fig, axes = plt.subplots(2, n_plots + 1, figsize=(2 * n_plots, 5)) + # Create plotly figure + fig = make_subplots( + rows=2, cols=n_plots + 1, + horizontal_spacing=0.01, + vertical_spacing=0.1 + ) - # Row labels - axes[0, 0].text(0.5, 0.5, 'Step-by-step', ha='center', va='center', fontsize=10) - axes[0, 0].axis('off') - axes[1, 0].text(0.5, 0.5, 'Direct', ha='center', va='center', fontsize=10) - axes[1, 0].axis('off') + # Add row labels + fig.add_annotation( + text='Step-by-step', xref="x domain", yref="y domain", + x=0.5, y=0.5, showarrow=False, font=dict(size=10), + row=1, col=1 + ) + fig.add_annotation( + text='Direct', xref="x domain", yref="y domain", + x=0.5, y=0.5, showarrow=False, font=dict(size=10), + row=2, col=1 + ) # Plot step-by-step diffusion for i, t in enumerate(timesteps): - axes[0, i + 1].imshow(tensor_to_image(xs[t]), cmap='gray') - axes[0, i + 1].set_title(f't={t}') - axes[0, i + 1].axis('off') + step_img = tensor_to_image(xs[t]).squeeze()[::-1] + fig.add_trace( + go.Heatmap(z=step_img, colorscale='gray', showscale=False), + row=1, col=i + 2 + ) + # Add title annotation for each column + fig.add_annotation( + text=f't={t}', + xref=f'x{i + 2} domain', yref=f'y{i + 2} domain', + x=0.5, y=1.15, showarrow=False, font=dict(size=10) + ) # Plot direct diffusion for comparison xt, _ = q_xt_x0(img, t) - axes[1, i + 1].imshow(tensor_to_image(xt), cmap='gray') - axes[1, i + 1].axis('off') + direct_img = tensor_to_image(xt).squeeze()[::-1] + fig.add_trace( + go.Heatmap(z=direct_img, colorscale='gray', showscale=False), + row=2, col=i + 2 + ) + + fig.update_layout( + title_text=f'Forward Diffusion Process (Label: {label})', + width=200 * n_plots, + height=500, + showlegend=False + ) + + # Hide axes for all subplots + fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False) + fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False) - fig.suptitle(f'Forward Diffusion Process (Label: {label})') - plt.tight_layout() + if save_path: + save_plot(fig, save_path) - if save_path is None: - save_path = os.path.join(PLOTS_DIR, 'forward_diffusion.png') - fig.savefig(save_path) - print(f"Saved forward diffusion visualization to {save_path}") - plt.close(fig) + return fig -def visualize_backward_diffusion(model: torch.nn.Module, save_path: str | None = None): +def visualize_backward_diffusion(model: torch.nn.Module, save_path: str | None = None) -> go.Figure: """ Visualize the backward (reverse) diffusion process. Args: model: Trained UNet model save_path: Path to save the visualization + Returns: + Plotly figure object """ ensure_dirs() model.eval() @@ -134,7 +168,8 @@ def visualize_backward_diffusion(model: torch.nn.Module, save_path: str | None = n_timesteps = 10 timesteps = np.linspace(1, DIFFU_STEPS, n_timesteps, dtype=int)[::-1] - fig, axes = plt.subplots(n_classes, n_timesteps, figsize=(2 * n_timesteps, 2 * n_classes)) + # Store images to plot later + images_to_plot = {} with torch.no_grad(): # Start from noise @@ -151,25 +186,60 @@ def visualize_backward_diffusion(model: torch.nn.Module, save_path: str | None = if t in timesteps: t_idx = timesteps.tolist().index(t) for class_idx in range(n_classes): - axes[class_idx, t_idx].imshow(tensor_to_image(x[class_idx]), cmap='gray') - if class_idx == 0: - axes[class_idx, t_idx].set_title(f't={t}') - if t_idx == 0: - axes[class_idx, t_idx].set_ylabel(f'Class {class_idx}') - axes[class_idx, t_idx].set_xticks([]) - axes[class_idx, t_idx].set_yticks([]) + images_to_plot[(class_idx, t_idx)] = { + 'img': tensor_to_image(x[class_idx]).squeeze()[::-1], + 't': t + } + + # Create plotly figure + fig = make_subplots( + rows=n_classes, cols=n_timesteps, + horizontal_spacing=0.01, + vertical_spacing=0.02 + ) + + for class_idx in range(n_classes): + for t_idx in range(n_timesteps): + data = images_to_plot[(class_idx, t_idx)] + fig.add_trace( + go.Heatmap(z=data['img'], colorscale='gray', showscale=False), + row=class_idx + 1, col=t_idx + 1 + ) + # Add column titles for first row + if class_idx == 0: + fig.add_annotation( + text=f"t={data['t']}", + xref=f'x{t_idx + 1} domain', yref=f'y{t_idx + 1} domain', + x=0.5, y=1.15, showarrow=False, font=dict(size=10) + ) + # Add row labels for first column + if t_idx == 0: + fig.add_annotation( + text=f'Class {class_idx}', + xref=f'x{class_idx * n_timesteps + 1} domain', + yref=f'y{class_idx * n_timesteps + 1} domain', + x=-0.2, y=0.5, showarrow=False, font=dict(size=10), + textangle=-90 + ) - fig.suptitle('Backward Diffusion Process') - plt.tight_layout() + fig.update_layout( + title_text='Backward Diffusion Process', + width=200 * n_timesteps, + height=200 * n_classes, + showlegend=False + ) - if save_path is None: - save_path = os.path.join(PLOTS_DIR, 'backward_diffusion.png') - fig.savefig(save_path) - print(f"Saved backward diffusion visualization to {save_path}") - plt.close(fig) + # Hide axes for all subplots + fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False) + fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False) + if save_path is not None: + save_plot(fig, save_path) -def generate_grid(model: torch.nn.Module, n_per_class: int = 10, save_path: str | None = None): + return fig + + +def generate_grid(model: torch.nn.Module, n_per_class: int = 10, save_path: str | None = None) -> go.Figure: """ Generate a grid of samples, organized by class. @@ -177,6 +247,8 @@ def generate_grid(model: torch.nn.Module, n_per_class: int = 10, save_path: str model: Trained UNet model n_per_class: Number of samples per class save_path: Path to save the grid + Returns: + Plotly figure object """ ensure_dirs() @@ -188,26 +260,45 @@ def generate_grid(model: torch.nn.Module, n_per_class: int = 10, save_path: str samples = generate_samples(model, labels=labels) samples = samples.cpu().numpy() - # Create grid - fig, axes = plt.subplots(NB_LABEL, n_per_class, figsize=(n_per_class, NB_LABEL)) + # Create plotly grid + fig = make_subplots( + rows=NB_LABEL, cols=n_per_class, + horizontal_spacing=0.01, + vertical_spacing=0.02 + ) for class_idx in range(NB_LABEL): for sample_idx in range(n_per_class): idx = class_idx * n_per_class + sample_idx - axes[class_idx, sample_idx].imshow(samples[idx].transpose(1, 2, 0).squeeze(), cmap='gray') - axes[class_idx, sample_idx].axis('off') - + img = samples[idx].transpose(1, 2, 0).squeeze()[::-1] + fig.add_trace( + go.Heatmap(z=img, colorscale='gray', showscale=False), + row=class_idx + 1, col=sample_idx + 1 + ) + # Add row labels for first column if sample_idx == 0: - axes[class_idx, sample_idx].set_ylabel(f'{class_idx}') + fig.add_annotation( + text=f'{class_idx}', + xref=f'x{class_idx * n_per_class + 1} domain', + yref=f'y{class_idx * n_per_class + 1} domain', + x=-0.2, y=0.5, showarrow=False, font=dict(size=10) + ) + + fig.update_layout( + title_text='Generated MNIST Digits', + width=100 * n_per_class, + height=100 * NB_LABEL, + showlegend=False + ) + + # Hide axes for all subplots + fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False) + fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False) - fig.suptitle('Generated MNIST Digits') - plt.tight_layout() + if save_path: + save_plot(fig, save_path) - if save_path is None: - save_path = os.path.join(PLOTS_DIR, 'generated_grid.png') - fig.savefig(save_path) - print(f"Saved generated grid to {save_path}") - plt.close(fig) + return fig if __name__ == "__main__": diff --git a/src/train_autoencoder.py b/src/train_autoencoder.py index 669ce47..6af1fc5 100644 --- a/src/train_autoencoder.py +++ b/src/train_autoencoder.py @@ -9,7 +9,8 @@ from src.config import DEVICE, BATCH_SIZE, NUM_WORKERS, CHECKPOINT_DIR, PLOTS_DI import os import torch import torch.nn as nn -import matplotlib.pyplot as plt +import plotly.graph_objects as go +from plotly.subplots import make_subplots from rich.progress import track import mlflow @@ -56,8 +57,9 @@ def train_autoencoder( # Setup training optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) - # Use MSE loss instead of BCE for better reconstruction of continuous values - criterion = nn.MSELoss() + # criterion = nn.BCELoss() + # criterion = nn.MSELoss() + criterion = nn.functional.binary_cross_entropy # Get dataloader dataloader = get_autoencoder_dataloader( @@ -80,9 +82,6 @@ def train_autoencoder( # Backward pass loss.backward() - # Gradient clipping to prevent exploding gradients - torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) - optimizer.step() epoch_loss += loss.item() @@ -94,14 +93,16 @@ def train_autoencoder( save_checkpoint(model, f"autoencoder_epoch_{epoch:03d}.pt") save_checkpoint(model, "autoencoder_latest.pt") - # Visualize results - visualize_reconstructions(model, dataloader) + # Visualize results + fig = visualize_reconstructions(model, dataloader) + + mlflow.log_figure(fig, f"reconstructions/epoch_{epoch:03d}.png") mlflow.end_run() return model -def visualize_reconstructions(model: AEModule, dataloader, n_samples: int = 10): +def visualize_reconstructions(model: AEModule, dataloader, n_samples: int = 10) -> go.Figure: """Visualize original, latent, and reconstructed images.""" model.eval() ensure_dirs() @@ -114,34 +115,56 @@ def visualize_reconstructions(model: AEModule, dataloader, n_samples: int = 10): latent = model.encode(x_batch) x_recon = model.decode(latent) - # Create visualization - fig, axes = plt.subplots(3, n_samples + 1, figsize=(2 * n_samples, 6)) + # Create visualization with plotly + fig = make_subplots( + rows=3, cols=n_samples + 1, + subplot_titles=[''] * (3 * (n_samples + 1)), + horizontal_spacing=0.01, + vertical_spacing=0.05 + ) - # Labels - axes[0, 0].text(0.5, 0.5, 'Original', ha='center', va='center', fontsize=12) - axes[0, 0].axis('off') - axes[1, 0].text(0.5, 0.5, 'Latent', ha='center', va='center', fontsize=12) - axes[1, 0].axis('off') - axes[2, 0].text(0.5, 0.5, 'Reconstructed', ha='center', va='center', fontsize=12) - axes[2, 0].axis('off') + # Add row labels as annotations + row_labels = ['Original', 'Latent', 'Reconstructed'] + for row_idx, label in enumerate(row_labels): + fig.add_annotation( + text=label, + xref="x domain", yref="y domain", + x=0.5, y=0.5, + showarrow=False, + font=dict(size=12), + row=row_idx + 1, col=1 + ) # Plot images for i in range(n_samples): - axes[0, i + 1].imshow(x_batch[i].cpu().squeeze().numpy(), cmap='gray') - axes[0, i + 1].axis('off') + # Original + fig.add_trace( + go.Heatmap(z=x_batch[i].cpu().squeeze().numpy()[::-1], colorscale='gray', showscale=False), + row=1, col=i + 2 + ) + # Latent + fig.add_trace( + go.Heatmap(z=latent[i].cpu().squeeze().numpy()[::-1], colorscale='gray', showscale=False), + row=2, col=i + 2 + ) + # Reconstructed + fig.add_trace( + go.Heatmap(z=x_recon[i].cpu().squeeze().numpy()[::-1], colorscale='gray', showscale=False), + row=3, col=i + 2 + ) - axes[1, i + 1].imshow(latent[i].cpu().squeeze().numpy(), cmap='gray') - axes[1, i + 1].axis('off') - - axes[2, i + 1].imshow(x_recon[i].cpu().squeeze().numpy(), cmap='gray') - axes[2, i + 1].axis('off') + fig.update_layout( + title_text='Autoencoder Results', + width=200 * n_samples, + height=600, + showlegend=False + ) - fig.suptitle('Autoencoder Results') - plt.tight_layout() + # Hide axes for all subplots + fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False) + fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False) - save_path = os.path.join(PLOTS_DIR, 'autoencoder_results.png') - fig.savefig(save_path) - plt.close(fig) + return fig if __name__ == "__main__": diff --git a/src/train_diffusion.py b/src/train_diffusion.py index cb70650..2ec4d60 100644 --- a/src/train_diffusion.py +++ b/src/train_diffusion.py @@ -18,7 +18,8 @@ import os import torch import torch.nn as nn import numpy as np -import matplotlib.pyplot as plt +import plotly.graph_objects as go +from plotly.subplots import make_subplots from rich.progress import track import mlflow @@ -156,20 +157,26 @@ def evaluate_and_log(model: nn.Module, epoch: int, predict_x0: bool = True): mlflow.log_metric("KL Divergence", kl_score, step=epoch) mlflow.log_metric("JSD", jsd_score, step=epoch) - # Log sample images - fig, axes = plt.subplots(4, 8, figsize=(16, 8)) - for i, ax in enumerate(axes.flat): - if i < len(fakes): - ax.imshow(fakes[i].transpose(1, 2, 0).squeeze(), cmap='gray') - ax.axis('off') - fig.suptitle(f"Generated Samples - Epoch {epoch}") - plt.tight_layout() + # Log sample images using plotly + fig = make_subplots(rows=4, cols=8, horizontal_spacing=0.01, vertical_spacing=0.02) + for i in range(min(32, len(fakes))): + row = i // 8 + 1 + col = i % 8 + 1 + img = fakes[i].transpose(1, 2, 0).squeeze()[::-1] + fig.add_trace( + go.Heatmap(z=img, colorscale='gray', showscale=False), + row=row, col=col + ) + fig.update_layout( + title_text=f"Generated Samples - Epoch {epoch}", + width=800, + height=400, + showlegend=False + ) + fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False) + fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False) - mlflow.log_figure(fig, f"samples_epoch_{epoch:03d}.png") - - # Save to plots folder - fig.savefig(os.path.join(PLOTS_DIR, f"samples_epoch_{epoch:03d}.png")) - plt.close(fig) + mlflow.log_figure(fig, f"samples/epoch_{epoch:03d}.png") if __name__ == "__main__": diff --git a/src/utils.py b/src/utils.py index f727a68..f34ec33 100644 --- a/src/utils.py +++ b/src/utils.py @@ -6,9 +6,9 @@ import os import io import numpy as np import torch -import scipy.linalg from PIL import Image -import matplotlib.pyplot as plt +import plotly.graph_objects as go +from plotly.subplots import make_subplots from .config import CHECKPOINT_DIR, PLOTS_DIR @@ -51,18 +51,18 @@ def tensor_to_images(tensor: torch.Tensor) -> np.ndarray: return img -def figure_to_image(figure: plt.Figure) -> np.ndarray: +def figure_to_image(figure: go.Figure) -> np.ndarray: """ - Convert a matplotlib figure to a numpy image array. + Convert a plotly figure to a numpy image array. Args: - figure: Matplotlib figure + figure: Plotly figure Returns: Numpy array of the figure image """ buf = io.BytesIO() - figure.savefig(buf, format='png') + figure.write_image(buf, format='png') buf.seek(0) image = np.array(Image.open(buf)) return image @@ -97,14 +97,14 @@ def load_checkpoint(model: torch.nn.Module, filename: str) -> torch.nn.Module: return model -def save_plot(figure: plt.Figure, filename: str): +def save_plot(figure: go.Figure, filename: str): """ - Save a matplotlib figure to the plots directory. + Save a plotly figure to the plots directory. Args: - figure: Matplotlib figure to save + figure: Plotly figure to save filename: Filename (will be saved in PLOTS_DIR) """ ensure_dirs() path = os.path.join(PLOTS_DIR, filename) - figure.savefig(path) + figure.write_image(path) |