init

app.py

@@ -7,9 +7,6 @@ import math
 import re
 from einops import rearrange
 from mmengine.config import Config
-from xtuner.registry import BUILDER
-from xtuner.model.utils import guess_load_checkpoint
-
 
 import matplotlib
 matplotlib.use("Agg")
@@ -18,6 +15,10 @@ import matplotlib.pyplot as plt
 from scripts.camera.cam_dataset import Cam_Generator
 from scripts.camera.visualization.visualize_batch import make_perspective_figures
 
+from mmengine.registry import Registry
+__all__ = ['BUILDER']
+BUILDER = Registry('builder')
+
 
 NUM = r"[+-]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][+-]?\d+)?"
 CAM_PATTERN = re.compile(r"(?:camera parameters.*?:|roll.*?:)\s*("+NUM+r")\s*,\s*("+NUM+r")\s*,\s*("+NUM+r")", re.IGNORECASE|re.DOTALL)
@@ -35,8 +36,8 @@ config = "configs/pipelines/stage_2_base.py"
 config = Config.fromfile(config)
 model = BUILDER.build(config.model).eval()
 checkpoint_path = "checkpoints/Puffin-Base.pth"
-state_dict = guess_load_checkpoint(checkpoint_path)
-model.load_state_dict(state_dict, strict=False)
+checkpoint = torch.load(checkpoint_path)
+info = model.load_state_dict(checkpoint, strict=False)
 
 if torch.cuda.is_available():
     model = model.to(torch.bfloat16).cuda()

configs/models/qwen2_5_1_5b_radio_sd3_dynamic_puffin.py

@@ -29,13 +29,13 @@ model = dict(type=Qwen2p5RadioStableDiffusion3HFDynamic,
                  hidden_size=1024,
                  intermediate_size=4096,
                  num_hidden_layers=6,
-                 _attn_implementation='flash_attention_2',
+                 #_attn_implementation='flash_attention_2',
                  num_attention_heads=16, ),
              connector_2=dict(
                  hidden_size=1024,
                  intermediate_size=4096,
                  num_hidden_layers=6,
-                 _attn_implementation='flash_attention_2',
+                 #_attn_implementation='flash_attention_2',
                  num_attention_heads=16, ),
              transformer=dict(
                  type=SD3Transformer2DModel.from_pretrained,
@@ -61,7 +61,7 @@ model = dict(type=Qwen2p5RadioStableDiffusion3HFDynamic,
                  type=AutoModelForCausalLM.from_pretrained,
                  pretrained_model_name_or_path=llm_name_or_path,
                  torch_dtype=torch.bfloat16,
-                 attn_implementation='flash_attention_2',
+                 #attn_implementation='flash_attention_2',
              ),
              tokenizer=dict(
                  type=AutoTokenizer.from_pretrained,