diff --git a/generate.py b/generate.py
index d06c6f2..f359e38 100644
--- a/generate.py
+++ b/generate.py
@@ -13,6 +13,8 @@ from byteplussdkarkruntime import Ark
 from django.conf import settings
 
 from item.models import Item
+from document.models import Document
+from archive.models import File, Stream
 
 
 os.environ["FAL_KEY"] = settings.FAL_KEY
@@ -479,7 +481,8 @@ def reshoot_item(item, extra_prompt=None, first_frame=None, keep=False):
     ai.data["seed"] = status["seed"]
     if first_frame:
         ai.data["firstframe"] = first_frame.split('?')[0]
-
+        if isinstance(first_frame, Document):
+            first_frame.add(ai)
     ai.save()
     if not keep:
         shutil.rmtree(os.path.dirname(output))
@@ -874,7 +877,6 @@ def luma_modify_item(item, prompt="", image_prompt=None, first_frame=None, keep=
 def add_ai_variant(item, video_path, type):
     if isinstance(item, str):
         item = Item.objects.get(public_id=item)
-    from archive.models import File, Stream
 
     ai = Item()
     ai.user = item.user
@@ -900,7 +902,6 @@ def add_ai_variant(item, video_path, type):
     return ai
 
 def add_ai_image(item, position, url, extension=None):
-    from document.models import Document
     if extension is None:
         extension = url.split('.')[-1].split('?')[0]
         if extension == 'jpeg': extension = 'jpg'