Upload folder using huggingface_hub

.idea/misc.xml

@@ -1,4 +1,7 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
+  <component name="Black">
+    <option name="sdkName" value="$USER_HOME$/envs/pt2.7" />
+  </component>
   <component name="ProjectRootManager" version="2" project-jdk-name="$USER_HOME$/envs/pt2.7" project-jdk-type="Python SDK" />
 </project>

.idea/workspace.xml

@@ -1,5 +1,8 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
+  <component name="AutoImportSettings">
+    <option name="autoReloadType" value="SELECTIVE" />
+  </component>
   <component name="ChangeListManager">
     <list default="true" id="9dd87dac-8a5e-4178-a1d7-afa664ac2f6a" name="Changes" comment="" />
     <option name="SHOW_DIALOG" value="false" />
@@ -7,6 +10,13 @@
     <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
     <option name="LAST_RESOLUTION" value="IGNORE" />
   </component>
+  <component name="FileTemplateManagerImpl">
+    <option name="RECENT_TEMPLATES">
+      <list>
+        <option value="Python Script" />
+      </list>
+    </option>
+  </component>
   <component name="ProjectColorInfo"><![CDATA[{
   "associatedIndex": 6
 }]]></component>
@@ -20,6 +30,10 @@
     "ModuleVcsDetector.initialDetectionPerformed": "true",
     "RunOnceActivity.ShowReadmeOnStart": "true",
     "last_opened_file_path": "/Users/wusize/projects/Puffin",
+    "node.js.detected.package.eslint": "true",
+    "node.js.detected.package.tslint": "true",
+    "node.js.selected.package.eslint": "(autodetect)",
+    "node.js.selected.package.tslint": "(autodetect)",
     "nodejs_package_manager_path": "npm",
     "settings.editor.selected.configurable": "com.jetbrains.python.configuration.PyActiveSdkModuleConfigurable",
     "vue.rearranger.settings.migration": "true"
@@ -40,11 +54,22 @@
       <option name="number" value="Default" />
       <option name="presentableId" value="Default" />
       <updated>1760056680813</updated>
-      <workItem from="1760056681869" duration="11000" />
+      <workItem from="1760056681869" duration="1047000" />
     </task>
     <servers />
   </component>
   <component name="TypeScriptGeneratedFilesManager">
     <option name="version" value="3" />
   </component>
+  <component name="XDebuggerManager">
+    <breakpoint-manager>
+      <breakpoints>
+        <line-breakpoint enabled="true" suspend="THREAD" type="python-line">
+          <url>file://$USER_HOME$/envs/pt2.7/lib/python3.10/site-packages/xtuner/dataset/map_fns/template_map_fn.py</url>
+          <line>1</line>
+          <option name="timeStamp" value="1" />
+        </line-breakpoint>
+      </breakpoints>
+    </breakpoint-manager>
+  </component>
 </project>

app.py

@@ -8,7 +8,7 @@ import math
 import re
 from einops import rearrange
 from mmengine.config import Config
-from xtuner.registry import BUILDER 
+from src.builder import BUILDER
 
 import matplotlib
 matplotlib.use("Agg")

requirements.txt

@@ -12,6 +12,3 @@ pillow==11.2.1
 scipy==1.15.2
 timm==0.9.12
 transformers==4.49.0
-xtuner==0.1.23
-deepspeed
-

src/builder.py

@@ -0,0 +1,4 @@
+from mmengine.registry import Registry
+__all__ = ['BUILDER']
+
+BUILDER = Registry('builder')

src/datasets/template_map_fn.py

@@ -0,0 +1,36 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from functools import partial
+
+from mmengine.utils.misc import get_object_from_string
+
+
+def template_map_fn(example, template):
+    conversation = example.get('conversation', [])
+    for i, single_turn_conversation in enumerate(conversation):
+        input = single_turn_conversation.get('input', '')
+        if input is None:
+            input = ''
+        input_text = template.INSTRUCTION.format(input=input, round=i + 1)
+        system = single_turn_conversation.get('system', '')
+        if system != '' and system is not None:
+            system = template.SYSTEM.format(system=system)
+            input_text = system + input_text
+        single_turn_conversation['input'] = input_text
+
+        if template.get('SUFFIX', None):
+            output_text = single_turn_conversation.get('output', '')
+            output_text += template.SUFFIX
+            single_turn_conversation['output'] = output_text
+
+        # SUFFIX_AS_EOS is False ==> need_eos_token is True
+        single_turn_conversation['need_eos_token'] = \
+            not template.get('SUFFIX_AS_EOS', False)
+        single_turn_conversation['sep'] = template.get('SEP', '')
+
+    return {'conversation': conversation}
+
+
+def template_map_fn_factory(template):
+    if isinstance(template, str):  # for resume
+        template = get_object_from_string(template)
+    return partial(template_map_fn, template=template)

src/datasets/utils.py

@@ -1,14 +1,42 @@
 import copy
 import random
-from xtuner.dataset.utils import get_bos_eos_token_ids
-from xtuner.utils import DEFAULT_IMAGE_TOKEN, IGNORE_INDEX, IMAGE_TOKEN_INDEX
 import json
 
+
+def get_bos_eos_token_ids(tokenizer):
+    if tokenizer.__class__.__name__ in [
+            'QWenTokenizer', 'QWen2Tokenizer', 'Qwen2TokenizerFast'
+    ]:
+        bos_token_id = []
+        eos_token_id = tokenizer.eos_token_id
+        assert eos_token_id is not None, \
+            'Please set eos_token for Qwen tokenizer!'
+    elif tokenizer.__class__.__name__ == 'ChatGLMTokenizer':
+        bos_token_id = [64790, 64792]
+        eos_token_id = tokenizer.eos_token_id
+    else:
+        bos_token_id = tokenizer.bos_token_id
+        eos_token_id = tokenizer.eos_token_id
+    if isinstance(bos_token_id, int):
+        bos_token_id = [bos_token_id]
+    if isinstance(eos_token_id, int):
+        eos_token_id = [eos_token_id]
+    return bos_token_id, eos_token_id
+
+
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN_INDEX = 0
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = '<image>'
+
 INPUT_IMAGE_TOKEN_INDEX = IMAGE_TOKEN_INDEX
 OUTPUT_IMAGE_TOKEN_INDEX = -300
 QUERY_TOKEN_INDEX = -400
 QUERY_TOKEN = '<query>'
 
+
+
+
 def crop2square(pil_img):
     width, height = pil_img.width, pil_img.height
 

src/models/puffin/model.py

@@ -11,16 +11,17 @@ from torch.autograd.function import Function
 from torch.nn.utils.rnn import pad_sequence
 from mmengine.logging import print_log
 from mmengine.model import BaseModel
-from xtuner.utils import IGNORE_INDEX
-from xtuner.registry import BUILDER
-from xtuner.model.utils import guess_load_checkpoint
-from xtuner.dataset.map_fns.template_map_fn import template_map_fn
+
+
+from src.builder import BUILDER
+from src.datasets.template_map_fn import template_map_fn
 from transformers.cache_utils import DynamicCache
 from diffusers.training_utils import compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3
 
 from src.models.connector import ConnectorConfig, ConnectorEncoder
 from src.models.stable_diffusion3.pipeline_stable_diffusion_3_dynamic import StableDiffusion3Pipeline
-from src.datasets.utils import encode_fn, QUERY_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, INPUT_IMAGE_TOKEN_INDEX
+from src.datasets.utils import (encode_fn, QUERY_TOKEN_INDEX, IGNORE_INDEX,
+                                DEFAULT_IMAGE_TOKEN, INPUT_IMAGE_TOKEN_INDEX)
 
 class _ScaleGradient(Function):
     @staticmethod
@@ -74,7 +75,7 @@ class Qwen2p5RadioStableDiffusion3HFDynamic(BaseModel):
                  fold_size=2,
                  unconditional=0.1,
                  unconditional_cross_view=0.1,
-                 pretrained_pth=None,
+                 # pretrained_pth=None,
                  use_activation_checkpointing=False,
                  *args, **kwargs):
         super().__init__()
@@ -136,11 +137,6 @@ class Qwen2p5RadioStableDiffusion3HFDynamic(BaseModel):
         if use_activation_checkpointing:
             self.llm.enable_input_require_grads()
             self.gradient_checkpointing_enable()
-
-        if pretrained_pth is not None:
-            pretrained_state_dict = guess_load_checkpoint(pretrained_pth)
-            info = self.load_state_dict(pretrained_state_dict, strict=False)
-            print_log(f'Load pretrained weight from {pretrained_pth}')
             
     @property
     def device(self):