""" CFG (Classifier-Free Guidance) — animation for @fminxyz Series 2, Post 5: "Как усилить кошачесть без классификатора" Visualization: Part 1 (frames 1-25): Vector diagram — ε_uncond, ε_cond, ε_cfg extrapolation Part 2 (frames 26-50): Sampling trajectories with w=1, 7, 30 """ import numpy as np import matplotlib.pyplot as plt import matplotlib.patches as mpatches from matplotlib.patches import FancyArrowPatch from matplotlib.animation import FuncAnimation, FFMpegWriter import matplotlib.patheffects as pe np.random.seed(42) # ---- Setup ---- FIG_SIZE = (10.8, 10.8) # 1080x1080 at 100 dpi DPI = 100 FPS = 2 TOTAL_FRAMES = 50 PART1_FRAMES = 25 PART2_FRAMES = 25 BG_DARK = "#0d1117" BG_PANEL = "#161b22" ACCENT_UNCOND = "#7c85ff" # blue — uncond ACCENT_COND = "#4cc9a0" # green — cond ACCENT_CFG = "#ff7043" # red/orange — cfg ACCENT_GRAY = "#8b949e" TEXT_COLOR = "#e6edf3" GRID_COLOR = "#21262d" # Two Gaussian modes: uncond (spread) and cond (focused on right) UNCOND_MEAN = np.array([0.0, 0.0]) COND_MEAN = np.array([2.5, 0.8]) def score_uncond(x, y): """Score of unconditional distribution (wider, centered at 0,0)""" sigma = 1.5 dx = UNCOND_MEAN[0] - x dy = UNCOND_MEAN[1] - y return np.array([dx / sigma**2, dy / sigma**2]) def score_cond(x, y): """Score of conditional distribution (focused on 2.5, 0.8)""" sigma = 0.9 dx = COND_MEAN[0] - x dy = COND_MEAN[1] - y return np.array([dx / sigma**2, dy / sigma**2]) def score_cfg(x, y, w): """CFG = uncond + w * (cond - uncond)""" su = score_uncond(x, y) sc = score_cond(x, y) return su + w * (sc - su) # ---- Trajectory simulation ---- def simulate_trajectory(x0, y0, w, n_steps=40, step_size=0.05): xs, ys = [x0], [y0] x, y = x0, y0 for _ in range(n_steps): s = score_cfg(x, y, w) noise = np.random.randn(2) * 0.05 if w < 2 else 0 x = x + step_size * s[0] + (noise[0] if isinstance(noise, np.ndarray) else 0) y = y + step_size * s[1] + (noise[1] if isinstance(noise, np.ndarray) else 0) xs.append(x) ys.append(y) return np.array(xs), np.array(ys) # Precompute trajectories np.random.seed(7) starts = [(-1.5, 1.0), (-2.0, -0.5), (-1.0, -1.5)] w_values = [1.0, 7.0, 30.0] trajectories = {} for w in w_values: trajs = [] for x0, y0 in starts: xs, ys = simulate_trajectory(x0, y0, w, n_steps=35, step_size=0.08) trajs.append((xs, ys)) trajectories[w] = trajs # Grid for background density xx, yy = np.meshgrid(np.linspace(-4, 5, 200), np.linspace(-3, 4, 200)) def density_uncond(x, y, sigma=1.5): return np.exp(-((x - UNCOND_MEAN[0])**2 + (y - UNCOND_MEAN[1])**2) / (2*sigma**2)) def density_cond(x, y, sigma=0.9): return np.exp(-((x - COND_MEAN[0])**2 + (y - COND_MEAN[1])**2) / (2*sigma**2)) # ---- Figure Setup ---- fig = plt.figure(figsize=FIG_SIZE, facecolor=BG_DARK, dpi=DPI) # Two panels side by side (split vertically) ax1 = fig.add_axes([0.02, 0.38, 0.96, 0.56]) # Vector diagram (top, taller) ax2 = fig.add_axes([0.02, 0.02, 0.96, 0.34]) # Trajectories (bottom) title_ax = fig.add_axes([0.02, 0.94, 0.96, 0.05]) title_ax.axis('off') for ax in [ax1, ax2]: ax.set_facecolor(BG_PANEL) for spine in ax.spines.values(): spine.set_color(GRID_COLOR) ax.tick_params(colors=ACCENT_GRAY, labelsize=9) ax.grid(True, color=GRID_COLOR, alpha=0.5, linewidth=0.5) # ---- Part 1: Vector diagram state ---- # Background density for ax1 Z_uncond = density_uncond(xx, yy, sigma=1.5) Z_cond = density_cond(xx, yy, sigma=0.9) Z_bg = 0.4 * Z_uncond + 0.7 * Z_cond ax1.contourf(xx, yy, Z_bg, levels=20, cmap='Blues', alpha=0.25) # Fixed point where we show vectors px, py = -0.8, 0.5 # Score vectors at this point su = score_uncond(px, py) * 0.6 sc = score_cond(px, py) * 0.6 # Labels for distributions ax1.text(UNCOND_MEAN[0] - 0.1, UNCOND_MEAN[1] + 1.8, 'p(x)\n(unconditional)', color=ACCENT_UNCOND, fontsize=9, ha='center', alpha=0.8) ax1.text(COND_MEAN[0] + 0.1, COND_MEAN[1] + 1.2, 'p(x|c)\n(кошка)', color=ACCENT_COND, fontsize=9, ha='center', alpha=0.8) # Draw distribution ellipses from matplotlib.patches import Ellipse e1 = Ellipse(UNCOND_MEAN, 5.5, 5.5, facecolor='none', edgecolor=ACCENT_UNCOND, linewidth=1.5, alpha=0.4, linestyle='--') e2 = Ellipse(COND_MEAN, 3.0, 3.0, facecolor='none', edgecolor=ACCENT_COND, linewidth=1.5, alpha=0.5) ax1.add_patch(e1) ax1.add_patch(e2) # Current point ax1.plot(px, py, 'o', color='white', markersize=10, zorder=10) ax1.text(px - 0.2, py - 0.4, 'xₜ', color=TEXT_COLOR, fontsize=12, fontweight='bold') # Static: ε_uncond vector arrow_uncond = ax1.annotate('', xy=(px + su[0], py + su[1]), xytext=(px, py), arrowprops=dict(arrowstyle='->', color=ACCENT_UNCOND, lw=2.5)) ax1.text(px + su[0] + 0.1, py + su[1] + 0.2, 'ε_uncond', color=ACCENT_UNCOND, fontsize=10, fontweight='bold') # Static: ε_cond vector arrow_cond = ax1.annotate('', xy=(px + sc[0], py + sc[1]), xytext=(px, py), arrowprops=dict(arrowstyle='->', color=ACCENT_COND, lw=2.5)) ax1.text(px + sc[0] + 0.05, py + sc[1] - 0.35, 'ε_cond', color=ACCENT_COND, fontsize=10, fontweight='bold') # Dynamic: ε_cfg (animated) cfg_arrow_line, = ax1.plot([], [], '-', color=ACCENT_CFG, lw=2.5, zorder=9) cfg_arrowhead = ax1.annotate('', xy=(px, py), xytext=(px, py), arrowprops=dict(arrowstyle='->', color=ACCENT_CFG, lw=2.5)) cfg_text = ax1.text(px, py, '', color=ACCENT_CFG, fontsize=10, fontweight='bold') w_label = ax1.text(-3.5, 3.2, '', color=TEXT_COLOR, fontsize=13, fontweight='bold', bbox=dict(boxstyle='round,pad=0.4', facecolor='#21262d', edgecolor=ACCENT_CFG, linewidth=1.5)) # Formula display formula_text = ax1.text(0.5, 0.02, 'ε_cfg = ε_uncond + w·(ε_cond − ε_uncond)', transform=ax1.transAxes, ha='center', va='bottom', color=TEXT_COLOR, fontsize=11, alpha=0.9, bbox=dict(boxstyle='round,pad=0.5', facecolor='#21262d', edgecolor=GRID_COLOR, linewidth=1)) ax1.set_xlim(-4, 5) ax1.set_ylim(-3, 4) ax1.set_title('Как работает Classifier-Free Guidance', color=TEXT_COLOR, fontsize=12, pad=8, fontweight='bold') # ---- Part 2: Trajectory panels ---- # Sub-axes for 3 w values side by side sub_width = 0.295 sub_positions = [ [0.05, 0.05, sub_width, 0.88], [0.37, 0.05, sub_width, 0.88], [0.69, 0.05, sub_width, 0.88], ] sub_axs = [] for i, pos in enumerate(sub_positions): sax = ax2.inset_axes(pos) sax.set_facecolor('#0d1117') for spine in sax.spines.values(): spine.set_color(GRID_COLOR) sax.set_xticks([]) sax.set_yticks([]) sub_axs.append(sax) # Background densities for sub-axes xx_s, yy_s = np.meshgrid(np.linspace(-3.5, 4.5, 100), np.linspace(-2.5, 3.5, 100)) Z_bg_s = 0.4 * density_uncond(xx_s, yy_s, sigma=1.5) + 0.7 * density_cond(xx_s, yy_s, sigma=0.9) w_colors = [ACCENT_UNCOND, ACCENT_COND, ACCENT_CFG] w_labels_text = ['w = 1\n(нет CFG)', 'w = 7\n(типичный)', 'w = 30\n(артефакты)'] for i, (sax, w, wc, wl) in enumerate(zip(sub_axs, w_values, w_colors, w_labels_text)): sax.contourf(xx_s, yy_s, Z_bg_s, levels=15, cmap='Blues', alpha=0.2) # Draw target modes sax.plot(*UNCOND_MEAN, 'o', color=ACCENT_UNCOND, markersize=8, alpha=0.5, zorder=5) sax.plot(*COND_MEAN, '*', color=ACCENT_COND, markersize=12, zorder=5) sax.text(COND_MEAN[0] + 0.1, COND_MEAN[1] + 0.6, '🐱', fontsize=12, ha='center') sax.set_xlim(-3.5, 4.5) sax.set_ylim(-2.5, 3.5) sax.set_title(wl, color=wc, fontsize=10, fontweight='bold', pad=4) # Trajectory artists (animated) traj_lines = [] traj_dots = [] for i, (sax, w, wc) in enumerate(zip(sub_axs, w_values, w_colors)): lines = [] dots = [] for j, (xs, ys) in enumerate(trajectories[w]): alpha = 0.6 + 0.2 * j line, = sax.plot([], [], '-', color=wc, linewidth=1.5, alpha=0.7, zorder=7) dot, = sax.plot([], [], 'o', color=wc, markersize=6, zorder=8) # Starting point sax.plot(xs[0], ys[0], 's', color='white', markersize=5, alpha=0.5, zorder=6) lines.append(line) dots.append(dot) traj_lines.append(lines) traj_dots.append(dots) ax2.axis('off') ax2.set_title('Эффект параметра w', color=TEXT_COLOR, fontsize=11, fontweight='bold', pad=0, x=0.5, y=0.95) # Title text title_ax.text(0.5, 0.5, 'Classifier-Free Guidance: усиливаем управление без классификатора', color=TEXT_COLOR, fontsize=13, fontweight='bold', ha='center', va='center', transform=title_ax.transAxes) # ---- Animation ---- max_traj_len = max(len(trajectories[w][0][0]) for w in w_values) def animate(frame): if frame < PART1_FRAMES: # Part 1: animate w from 0.5 to 30 t = frame / (PART1_FRAMES - 1) # Non-linear: slow at start, fast in middle, slow at end w_anim = np.exp(t * np.log(31)) * 0.5 # 0.5 → ~15 if frame > 20: w_anim = 1 + (frame - 1) * 2.0 # quickly go to 30 w_anim = min(w_anim, 30.0) sc_cfg = score_cfg(px, py, w_anim) * 0.6 cfg_end = (px + sc_cfg[0], py + sc_cfg[1]) # Update CFG arrow cfg_arrow_line.set_data([px, cfg_end[0]], [py, cfg_end[1]]) cfg_arrowhead.remove() new_arrow = ax1.annotate('', xy=cfg_end, xytext=(px, py), arrowprops=dict(arrowstyle='->', color=ACCENT_CFG, lw=2.5)) globals()['cfg_arrowhead_current'] = new_arrow cfg_text.set_position((cfg_end[0] + 0.1, cfg_end[1] + 0.1)) cfg_text.set_text(f'ε_cfg') w_label.set_text(f'w = {w_anim:.0f}') # Keep trajectory panels static in part 1 for i, w_traj in enumerate(w_values): n_show = min(10, max_traj_len) for j, (line, dot) in enumerate(zip(traj_lines[i], traj_dots[i])): xs, ys = trajectories[w_traj][j] line.set_data(xs[:n_show], ys[:n_show]) dot.set_data([xs[min(n_show-1, len(xs)-1)]], [ys[min(n_show-1, len(ys)-1)]]) else: # Part 2: animate trajectories t2_frame = frame - PART1_FRAMES n_show = int(1 + t2_frame * (max_traj_len - 1) / (PART2_FRAMES - 1)) n_show = min(n_show, max_traj_len) for i, w_traj in enumerate(w_values): for j, (line, dot) in enumerate(zip(traj_lines[i], traj_dots[i])): xs, ys = trajectories[w_traj][j] end_idx = min(n_show, len(xs)) line.set_data(xs[:end_idx], ys[:end_idx]) dot.set_data([xs[end_idx-1]], [ys[end_idx-1]]) # Keep CFG vector at w=7 in part 2 sc_cfg = score_cfg(px, py, 7.0) * 0.6 cfg_end = (px + sc_cfg[0], py + sc_cfg[1]) cfg_arrow_line.set_data([px, cfg_end[0]], [py, cfg_end[1]]) cfg_text.set_position((cfg_end[0] + 0.1, cfg_end[1] + 0.1)) cfg_text.set_text('ε_cfg (w=7)') w_label.set_text('w = 7') # Remove the buggy globals approach, simplify animation def make_cfg_arrow_dynamic(): """Simpler approach: just remove and re-add arrow each frame is too slow. Instead, use a line + triangle patch.""" pass # Actually, the annotate approach won't work well for animation. # Let's use a simpler quiver approach instead. # Redo ax1 with quiver ax1.cla() ax1.set_facecolor(BG_PANEL) for spine in ax1.spines.values(): spine.set_color(GRID_COLOR) ax1.tick_params(colors=ACCENT_GRAY, labelsize=9) ax1.grid(True, color=GRID_COLOR, alpha=0.5, linewidth=0.5) ax1.contourf(xx, yy, Z_bg, levels=20, cmap='Blues', alpha=0.25) e1 = Ellipse(UNCOND_MEAN, 5.5, 5.5, facecolor='none', edgecolor=ACCENT_UNCOND, linewidth=1.5, alpha=0.4, linestyle='--') e2 = Ellipse(COND_MEAN, 3.0, 3.0, facecolor='none', edgecolor=ACCENT_COND, linewidth=1.5, alpha=0.5) ax1.add_patch(e1) ax1.add_patch(e2) ax1.text(UNCOND_MEAN[0] - 0.1, UNCOND_MEAN[1] + 2.0, 'p(x)\nunconditional', color=ACCENT_UNCOND, fontsize=9, ha='center', alpha=0.8) ax1.text(COND_MEAN[0] + 0.3, COND_MEAN[1] + 1.4, 'p(x|c)\nкошка 🐱', color=ACCENT_COND, fontsize=9, ha='center', alpha=0.9) ax1.plot(px, py, 'o', color='white', markersize=10, zorder=10) ax1.text(px - 0.35, py - 0.45, 'xₜ', color=TEXT_COLOR, fontsize=12, fontweight='bold') # Static ε_uncond ax1.quiver(px, py, su[0], su[1], color=ACCENT_UNCOND, scale=1, scale_units='xy', angles='xy', width=0.008, headwidth=4, headlength=4, zorder=8) ax1.text(px + su[0] + 0.15, py + su[1] + 0.15, 'ε_uncond', color=ACCENT_UNCOND, fontsize=10, fontweight='bold', zorder=9) # Static ε_cond ax1.quiver(px, py, sc[0], sc[1], color=ACCENT_COND, scale=1, scale_units='xy', angles='xy', width=0.008, headwidth=4, headlength=4, zorder=8) ax1.text(px + sc[0] + 0.05, py + sc[1] - 0.4, 'ε_cond', color=ACCENT_COND, fontsize=10, fontweight='bold', zorder=9) # Dynamic ε_cfg via quiver (we'll update U, V) cfg_quiver = ax1.quiver(px, py, 0.01, 0.01, color=ACCENT_CFG, scale=1, scale_units='xy', angles='xy', width=0.012, headwidth=4.5, headlength=5, zorder=9) cfg_label = ax1.text(px, py, '', color=ACCENT_CFG, fontsize=10, fontweight='bold', zorder=10) w_display = ax1.text(-3.5, 3.3, 'w = 1', color=TEXT_COLOR, fontsize=14, fontweight='bold', bbox=dict(boxstyle='round,pad=0.5', facecolor='#21262d', edgecolor=ACCENT_CFG, linewidth=2), zorder=11) ax1.text(0.5, 0.02, 'ε_cfg = ε_uncond + w · (ε_cond − ε_uncond)', transform=ax1.transAxes, ha='center', va='bottom', color=TEXT_COLOR, fontsize=11, alpha=0.9, bbox=dict(boxstyle='round,pad=0.5', facecolor='#21262d', edgecolor=GRID_COLOR, linewidth=1)) ax1.set_xlim(-4, 5) ax1.set_ylim(-3, 4) ax1.set_title('Как работает Classifier-Free Guidance', color=TEXT_COLOR, fontsize=12, pad=8, fontweight='bold') # W schedule for part 1: 1→7→30 def w_schedule(frame): """Returns w value for given frame in part 1 (0..24)""" t = frame / (PART1_FRAMES - 1) # 0..1 if t < 0.5: return 1.0 + (7.0 - 1.0) * (t / 0.5) # 1→7 else: return 7.0 + (30.0 - 7.0) * ((t - 0.5) / 0.5) # 7→30 def animate2(frame): artists = [] if frame < PART1_FRAMES: w_cur = w_schedule(frame) sc_cfg = score_cfg(px, py, w_cur) * 0.6 cfg_quiver.set_UVC(sc_cfg[0], sc_cfg[1]) cfg_ex, cfg_ey = px + sc_cfg[0], py + sc_cfg[1] cfg_label.set_position((cfg_ex + 0.1, cfg_ey + 0.15)) cfg_label.set_text('ε_cfg') w_display.set_text(f'w = {w_cur:.0f}') # Trajectory: show first 8 steps only for i, w_traj in enumerate(w_values): for j, (line, dot) in enumerate(zip(traj_lines[i], traj_dots[i])): xs, ys = trajectories[w_traj][j] n_show = 8 line.set_data(xs[:n_show], ys[:n_show]) dot.set_data([xs[n_show-1]], [ys[n_show-1]]) else: # Part 2: animate trajectories growing t2_frame = frame - PART1_FRAMES n_show = int(1 + t2_frame * (max_traj_len - 1) / max(PART2_FRAMES - 1, 1)) n_show = min(n_show, max_traj_len) for i, w_traj in enumerate(w_values): for j, (line, dot) in enumerate(zip(traj_lines[i], traj_dots[i])): xs, ys = trajectories[w_traj][j] end_idx = min(n_show, len(xs)) line.set_data(xs[:end_idx], ys[:end_idx]) dot.set_data([xs[end_idx-1]], [ys[end_idx-1]]) # Keep cfg arrow at w=30 for part 2 w_cur = 30.0 sc_cfg = score_cfg(px, py, w_cur) * 0.6 cfg_quiver.set_UVC(sc_cfg[0], sc_cfg[1]) cfg_ex, cfg_ey = px + sc_cfg[0], py + sc_cfg[1] cfg_label.set_position((cfg_ex + 0.1, cfg_ey + 0.15)) cfg_label.set_text('ε_cfg (w=30)') w_display.set_text('w = 30') # Create animation ani = FuncAnimation(fig, animate2, frames=TOTAL_FRAMES, interval=500, blit=False) # Save output_path = '/root/Strategy/content/drafts/cfg_animation.mp4' writer = FFMpegWriter(fps=FPS, metadata={'title': 'CFG Animation'}, extra_args=['-vcodec', 'libx264', '-pix_fmt', 'yuv420p', '-crf', '23']) ani.save(output_path, writer=writer, dpi=DPI) print(f"Saved: {output_path}") # Verify import os size_mb = os.path.getsize(output_path) / 1024 / 1024 print(f"File size: {size_mb:.2f} MB") print("Done!")