update code

app.py

@@ -446,7 +446,7 @@ def duplicate_batch(tensor, batch_size=2):
 
 @torch.no_grad()
 @spaces.GPU(duration=200)
-def avatar_generation(items, save_path_base, video_path_input, source_type, is_styled, styled_img):
+def avatar_generation(items, save_path_base, video_path_input, source_type, is_styled, styled_img, image_name_true):
     """
     Generate avatars from input images.
 
@@ -480,8 +480,6 @@ def avatar_generation(items, save_path_base, video_path_input, source_type, is_s
     # ws_avg.to(device)
     # DiT_model.to(device)
     # Set up face verse for amimation
-
-
     if source_type == 'example':
         input_img_fvid = './demo_data/source_img/img_generate_different_domain/coeffs/demo_imgs'
         input_img_motion = './demo_data/source_img/img_generate_different_domain/motions/demo_imgs'
@@ -502,8 +500,17 @@ def avatar_generation(items, save_path_base, video_path_input, source_type, is_s
             raise ValueError("Batch size > 1 not implemented")
 
         image_dir = chunk[0]
+        image_name = os.path.splitext(image_name_true)[0]
+
+        # # image_name = os.path.splitext(os.path.basename(image_dir))[0]
+        # if source_type == 'custom':
+        #     image_name = os.path.splitext(image_name_true)[0]
+        # else:
+        #     image_name = os.path.splitext(os.path.basename(image_dir))[0]
+
+
+
 
-        image_name = os.path.splitext(os.path.basename(image_dir))[0]
         dino_img, clip_image = image_process(image_dir, clip_image_processor, dino_img_processor, device)
 
         clip_feature = image_encoder(clip_image, output_hidden_states=True).hidden_states[-2]
@@ -606,8 +613,8 @@ def avatar_generation(items, save_path_base, video_path_input, source_type, is_s
             Image.fromarray(final_out_show, 'RGB').save(os.path.join(save_frames_path_outshow, frame_name))
 
         # Generate videos
-        images_to_video(save_frames_path_out, os.path.join(save_path_base, image_name + '_out.mp4'))
-        images_to_video(save_frames_path_depth, os.path.join(save_path_base, image_name + '_depth.mp4'))
+        images_to_video(save_frames_path_out, os.path.join(save_path_base, image_name + video_name+ '_out.mp4'))
+        images_to_video(save_frames_path_depth, os.path.join(save_path_base, image_name + video_name+ '_out.mp4'))
 
         logging.info(f"✅ Video generation completed successfully!")
         return os.path.join(save_path_base, image_name + video_name+ '_out.mp4'),  os.path.join(save_path_base, image_name + video_name+'_depth.mp4')
@@ -624,30 +631,31 @@ def assert_input_image(input_image):
         raise gr.Error("No image selected or uploaded!")
 
 @spaces.GPU(duration=100)
-def process_image(input_image, source_type, is_style, save_dir):
+def process_image(input_image_dir, source_type, is_style, save_dir):
     """ 🎯 处理 input_image，根据是否是示例图片执行不同逻辑 """
     process_img_input_dir = os.path.join(save_dir, 'input_image')
     process_img_save_dir = os.path.join(save_dir, 'processed_img')
+    image_name_true = os.path.basename(input_image_dir)
     os.makedirs(process_img_save_dir, exist_ok=True)
     os.makedirs(process_img_input_dir, exist_ok=True)
     if source_type == "example":
-        return input_image, source_type
+        input_image = Image.open(input_image)
+        return input_image, source_type, image_name_true
     else:
         # input_process_model.inference(input_image, process_img_save_dir)
-        shutil.copy(input_image, process_img_input_dir)
+        shutil.copy(input_image_dir, process_img_input_dir)
         input_process_model.inference(process_img_input_dir, process_img_save_dir, is_img=True, is_video=False)
-        img_name = os.path.basename(input_image)
-        imge_dir = os.path.join(save_dir, 'processed_img/dataset/images512x512/input_image', img_name)
-        return imge_dir, source_type  # 这里替换成 处理用户上传图片的逻辑
+        imge_dir = os.path.join(save_dir, 'processed_img/dataset/images512x512/input_image', image_name_true)
+        image = Image.open(imge_dir)
+        return image, source_type, image_name_true  # 这里替换成 处理用户上传图片的逻辑
 
 @spaces.GPU(duration=100)
 @torch.no_grad()
-def style_transfer(processed_image, style_prompt, cfg, strength, save_base):
+def style_transfer(processed_image, style_prompt, cfg, strength, save_base,image_name_true):
     """
     🎭 这个函数用于风格转换
     ✅ 你可以在这里填入你的风格化代码
     """
-    pipeline_sd.to(device)
     src_img_pil = Image.open(processed_image)
     img_name = os.path.basename(processed_image)
     save_dir = os.path.join(save_base, 'style_img')
@@ -663,8 +671,8 @@ def style_transfer(processed_image, style_prompt, cfg, strength, save_base):
         num_inference_steps=30,
         controlnet_conditioning_scale=1.5
     )['images'][0]
-    trg_img_pil.save(os.path.join(save_dir, img_name))
-    return os.path.join(save_dir, img_name) # 🚨 这里需要替换成你的风格转换逻辑
+    trg_img_pil.save(os.path.join(save_dir, image_name_true))
+    return trg_img_pil # 🚨 这里需要替换成你的风格转换逻辑
 
 
 def reset_flag():
@@ -808,6 +816,8 @@ def launch_gradio_app():
         is_from_example = gr.State(value=True)
         is_styled = gr.State(value=False)
         working_dir = gr.State()
+        image_name_true = gr.State()
+
 
         with gr.Row():
             with gr.Column(variant='panel'):
@@ -932,13 +942,14 @@ def launch_gradio_app():
                             format="mp4", height=512, width=512,
                             autoplay=True
                         )
-        def apply_style_and_mark(processed_image, style_choice, cfg, strength, working_dir):
-            styled = style_transfer(processed_image, styles[style_choice], cfg, strength, working_dir)
+        def apply_style_and_mark(processed_image, style_choice, cfg, strength, working_dir, image_name_true):
+            styled = style_transfer(processed_image, styles[style_choice], cfg, strength, working_dir, image_name_true)
             return styled, True
 
         def process_image_and_enable_style(input_image, source_type, is_styled, wd):
-            processed_result, updated_source_type = process_image(input_image, source_type, is_styled, wd)
-            return processed_result, updated_source_type, gr.update(interactive=True), gr.update(interactive=True)
+            processed_result, updated_source_type, image_name_true = process_image(input_image, source_type, is_styled, wd)
+            return processed_result, updated_source_type, gr.update(interactive=True), gr.update(interactive=True), image_name_true
+        
         processed_image_button.click(
             fn=prepare_working_dir,
             inputs=[working_dir, is_styled],
@@ -947,17 +958,17 @@ def launch_gradio_app():
         ).success(
             fn=process_image_and_enable_style,
             inputs=[input_image, source_type, is_styled, working_dir],
-            outputs=[processed_image, source_type, style_button, submit],
+            outputs=[processed_image, source_type, style_button, submit, image_name_true],
             queue=True
         )
         style_button.click(
             fn=apply_style_and_mark,
-            inputs=[processed_image, style_choice, cfg_slider, strength_slider, working_dir],
+            inputs=[processed_image, style_choice, cfg_slider, strength_slider, working_dir, image_name_true],
             outputs=[style_image, is_styled]
         )
         submit.click(
             fn=avatar_generation,
-            inputs=[processed_image, working_dir, video_input, source_type, is_styled, style_image],
+            inputs=[processed_image, working_dir, video_input, source_type, is_styled, style_image, image_name_true],
             outputs=[output_video, output_video_1],  # ⏳ 稍后展示视频
             queue=True
         )