Migrate torch.cuda.amp.autocast to torch.amp

[xyf5432]

Fixes #592

Summary

torch.cuda.amp.autocast has been deprecated since torch 2.4 and is scheduled for removal. This PR migrates the seven uses in the core flagai/ package to torch.amp without raising the torch floor, since the repo documents PyTorch >= 1.8.0 (README "Requirements and Installation") and pins no torch version:

• Imports (6 lines, all switching to a fallback import — torch.amp on torch >= 2.0, torch.cuda.amp below, where it is not yet deprecated):

  try:  # torch >= 2.0
      from torch.amp import autocast
  except ImportError:  # torch < 2.0
      from torch.cuda.amp import autocast

  • flagai/model/predictor/predictor.py:21 (public Predictor API)
  • flagai/model/mm/AltDiffusion.py:17, AltDiffusionM18.py:17, 19, Unets/Unet.py:19, modules/diffusionmodules/openaimodel.py:10
• Call sites (5, in the same files): with autocast(): → with _autocast():, where a version-guarded helper keeps both call signatures correct:

  def _autocast(**kwargs):
      """torch.amp.autocast("cuda") on torch >= 2.0, else torch.cuda.amp.autocast."""
      if hasattr(torch, "amp"):
          return torch.amp.autocast("cuda", **kwargs)
      return torch.cuda.amp.autocast(**kwargs)

  Why not a plain autocast("cuda")? The two APIs take different call forms: torch.amp.autocast (>= 2.0) requires the device type as the first positional argument, while torch.cuda.amp.autocast (< 2.0) has enabled as its first positional parameter — a bare autocast("cuda") would silently feed "cuda" into the old API's enabled slot. That happens to behave equivalently (truthy check), but it relies on unspecified behavior, so the helper dispatches to the exact signature each version expects.
• modules/diffusionmodules/util.py:136 — same helper pattern for the kwargs form inside the autograd function (torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs) takes keyword args only):

  def _autocast_ctx(**kwargs):
      """torch.amp.autocast("cuda") on torch >= 2.0, else torch.cuda.amp.autocast."""
      if hasattr(torch, "amp"):
          return torch.amp.autocast("cuda", **kwargs)
      return torch.cuda.amp.autocast(**kwargs)

The vendored examples/ca-lora/src/section-4.1/opendelta/ copy is third-party embedded code and intentionally left untouched.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Validation

Verified on torch 2.11.0 with warnings.simplefilter("error", FutureWarning):

• _autocast() / _autocast_ctx() on torch >= 2.0 dispatch to torch.amp.autocast("cuda", ...) — including the enabled=... kwargs form — and run warning-free.
• Fallback branch reachability: with torch.amp masked (simulating torch < 2.0), the ImportError is raised, the imports fall through to torch.cuda.amp.autocast, and both helpers dispatch to torch.cuda.amp.autocast(**kwargs) — the signature the old API actually takes (the old API emits no warning before torch 2.4).
• Control group: torch.cuda.amp.autocast() raises FutureWarning under the same filter — the fix is effective and the filter is sensitive.
• py_compile passes on all six touched files; the core flagai/ package has zero bare torch.cuda.amp references (grep-verified).

User impact

No behavior change on any supported torch version (the documented floor of 1.8 is preserved); the FutureWarning emitted by any Predictor-based inference and AltDiffusion inference/training on torch 2.4+ is eliminated, as is the breakage risk once torch.cuda.amp is removed.

Notes for reviewers

• All five call sites previously used the bare with autocast(): form (default enabled, float16), and _autocast() preserves exactly that on both torch paths — no semantics change.
• Same migration as ultralytics/yolov5#13244 and huggingface/lerobot#3167.