test

app.py

@@ -400,36 +400,27 @@ class DreamblendGUI:
             canvas_size=400
         ), draggable_img
 
-    def save_image(self, save_path = "/mnt/bn/hjj-humanseg-lq/SubjectDriven/DreamFuse/debug"):
-        global generated_images
-        save_name = self.get_next_sequence(save_path)
-        generated_images[0].save(os.path.join(save_path, f"{save_name}_0_ori.png"))
-        generated_images[1].save(os.path.join(save_path, f"{save_name}_0.png"))
-        generated_images[2].save(os.path.join(save_path, f"{save_name}_1.png"))
-        generated_images[3].save(os.path.join(save_path, f"{save_name}_2.png"))
-        generated_images[4].save(os.path.join(save_path, f"{save_name}_0_mask.png"))
-        generated_images[5].save(os.path.join(save_path, f"{save_name}_0_mask_scale.png"))
-        generated_images[6].save(os.path.join(save_path, f"{save_name}_0_scale.png"))
-        generated_images[7].save(os.path.join(save_path, f"{save_name}_2_pasted.png"))
-        
 
     def create_gui(self):
         config = InferenceConfig()
         config.lora_id = 'LL3RD/DreamFuse'
         
-        pipeline = DreamFuseInference(config)
-        pipeline.gradio_generate = spaces.GPU(duratioin=120)(pipeline.gradio_generate)
+        pipeline = None
+        # pipeline = DreamFuseInference(config)
+        # pipeline.gradio_generate = spaces.GPU(duratioin=120)(pipeline.gradio_generate)
         """创建 Gradio 界面"""
         with gr.Blocks(css=self.css_style) as demo:
             modified_fg_state = gr.State() 
-            gr.Markdown("# Dreamblend-GUI-dirtydata")
-            gr.Markdown("通过上传背景图与前景图生成带有可拖拽/缩放预览的合成图像，同时支持 Seed 设置和 Prompt 文本输入。")
+            gr.Markdown("# DreamFuse: 3 Easy Steps to Create Your Fusion Image")
+            gr.Markdown("1. Upload the foreground and background images you want to blend.")
+            gr.Markdown("2. Click 'Generate Canvas' to preview the result. You can then drag and resize the foreground object to position it as you like.")
+            gr.Markdown("3. Click 'Run Model' to create the final fused image.")
             with gr.Row():
                 with gr.Column(scale=1):
-                    gr.Markdown("### 上传图片")
-                    background_img_in = gr.Image(label="背景图片", type="pil", height=240, width=240)
-                    draggable_img_in = gr.Image(label="前景图片", type="pil", image_mode="RGBA", height=240, width=240)
-                    generate_btn = gr.Button("生成可拖拽画布")
+                    gr.Markdown("### FG&BG Image Upload")
+                    background_img_in = gr.Image(label="Background Image", type="pil", height=240, width=240)
+                    draggable_img_in = gr.Image(label="Foreground Image", type="pil", image_mode="RGBA", height=240, width=240)
+                    generate_btn = gr.Button("Generate Canvas")
 
                     with gr.Row():
                         gr.Examples(
@@ -438,39 +429,38 @@ class DreamblendGUI:
                             elem_id="small-examples"
                         )
                 with gr.Column(scale=1):
-                    gr.Markdown("### 预览区域")
-                    html_out = gr.HTML(label="预览与拖拽", elem_id="canvas_preview")
+                    gr.Markdown("### Preview Region")
+                    html_out = gr.HTML(label="drag and resize", elem_id="canvas_preview")
 
             with gr.Row():
                 with gr.Column(scale=1):
-                    gr.Markdown("### 参数设置")
-                    seed_slider = gr.Slider(minimum=0, maximum=10000, step=1, label="Seed", value=42)
+                    gr.Markdown("### Parameters")
+                    seed_slider = gr.Slider(minimum=-1, maximum=100000, step=1, label="Seed", value=12345)
                     cfg_slider = gr.Slider(minimum=1, maximum=10, step=0.1, label="CFG", value=3.5)
                     size_select = gr.Radio(
                         choices=["512", "768", "1024"],
                         value="512",
                         label="生成质量(512-差 1024-好)",
                     )
-                    prompt_text = gr.Textbox(label="Prompt", placeholder="输入文本提示", value="")
-                    text_strength = gr.Slider(minimum=1, maximum=10, step=1, label="Text Strength", value=1)
-                    enable_gui = gr.Checkbox(label="启用GUI", value=True)
-                    enable_truecfg = gr.Checkbox(label="启用TrueCFG", value=False)
-                    enable_save = gr.Button("保存图片 (内部测试)", visible=True)
+                    prompt_text = gr.Textbox(label="Prompt", placeholder="text prompt", value="")
+                    text_strength = gr.Slider(minimum=1, maximum=10, step=1, label="Text Strength", value=1, visible=False)
+                    enable_gui = gr.Checkbox(label="启用GUI", value=True, visible=False)
+                    enable_truecfg = gr.Checkbox(label="TrueCFG", value=False, visible=False)
                 with gr.Column(scale=1):
-                    gr.Markdown("### 模型生成结果")
-                    model_generate_btn = gr.Button("模型生成")
+                    gr.Markdown("### Model Result")
+                    model_generate_btn = gr.Button("Run Model")
                     transformation_text = gr.Textbox(label="Transformation Info", elem_id="transformation_info", visible=False)
-                    model_output = gr.Image(label="模型输出", type="pil")
+                    model_output = gr.Image(label="Model Output", type="pil")
             
-            # 交互事件绑定
-            enable_save.click(fn=self.save_image, inputs=None, outputs=None)
+
             generate_btn.click(
                 fn=self.on_upload,
                 inputs=[background_img_in, draggable_img_in],
                 outputs=[html_out, modified_fg_state],
             )
             model_generate_btn.click(
-                fn=pipeline.gradio_generate,
+                # fn=pipeline.gradio_generate,
+                fn=self.pil_to_base64,
                 inputs=[background_img_in, modified_fg_state, transformation_text, seed_slider, \
                     prompt_text, enable_gui, cfg_slider, size_select, text_strength, enable_truecfg],
                 outputs=model_output

dreamfuse_inference.py

@@ -168,7 +168,6 @@ def make_image_grid(images, rows, cols, size=None):
 class DreamFuseInference:
     def __init__(self, config: InferenceConfig):
         self.config = config
-        print(config.device)
         self.device = torch.device(config.device)
         torch.backends.cuda.matmul.allow_tf32 = True
         seed_everything(config.seed)
@@ -348,16 +347,15 @@ class DreamFuseInference:
 
     @torch.inference_mode()
     def gradio_generate(self, background_img, foreground_img, transformation_info, seed, prompt, enable_gui, cfg=3.5, size_select="1024", text_strength=1, truecfg=False):
-        print("!"*10)
-        """使用 DreamFuseInference 进行模型推理"""
         try:
             trans = json.loads(transformation_info)
         except:
             trans = {}
 
         size_select = int(size_select)
+        if size_select == 1024: text_strength = 5
+        if size_select == 768: text_strength = 3
         
-        # import pdb; pdb.set_trace()
         r, g, b, ori_a = foreground_img.split()
         fg_img_scale, fg_img = self.transform_foreground_original(foreground_img, background_img, trans)
 
@@ -370,9 +368,7 @@ class DreamFuseInference:
         ori_a = ori_a.convert("L")
         new_a = new_a.convert("L")
         foreground_img.paste((255, 255, 255), mask=ImageOps.invert(ori_a))
-        print("0"*10)
-        print(foreground_img.size)
-        print(background_img.size)
+
         images = self.model_generate(foreground_img.copy(), background_img.copy(),
                         ori_a, new_a,
                         enable_mask_affine=enable_gui, 
@@ -386,16 +382,15 @@ class DreamFuseInference:
         images = Image.fromarray(images[0], "RGB")
 
         images = images.resize(background_img.size)
-        images_save = images.copy()
+        # images_save = images.copy()
             
-        images.thumbnail((640, 640), Image.LANCZOS)
+        # images.thumbnail((640, 640), Image.LANCZOS)
         return images
 
 
     @torch.inference_mode()
     def model_generate(self, fg_image, bg_image, ori_fg_mask, new_fg_mask, enable_mask_affine=True, prompt="", offset_cond=None, seed=None, cfg=3.5, size_select=1024, text_strength=1, truecfg=False):
         batch_size = 1
-        print("-3"*10)
         # Prepare images
         # adjust bg->fg size
         fg_image, ori_fg_mask = adjust_fg_to_bg(fg_image, ori_fg_mask, bg_image.size)
@@ -410,7 +405,6 @@ class DreamFuseInference:
             new_fg_mask = new_fg_mask.resize(bucket_size)
             mask_affine = get_mask_affine(new_fg_mask, ori_fg_mask)
 
-        print("-2"*10)
         # Get embeddings
         prompt_embeds, pooled_prompt_embeds, text_ids = self._compute_text_embeddings(prompt)
 
@@ -428,7 +422,6 @@ class DreamFuseInference:
         if seed is None:
             seed = self.config.seed
         generator = torch.Generator(device=self.device).manual_seed(seed)
-        print("-1"*10)
         # Prepare condition latents
         condition_image_latents = self._encode_images([fg_image, bg_image])
         
@@ -445,7 +438,6 @@ class DreamFuseInference:
                 )
             )
 
-        print(1)
         if mask_affine is not None:
             affine_H, affine_W = condition_image_latents.shape[2] // 2, condition_image_latents.shape[3] // 2
             scale_factor = 1 / 16
@@ -457,7 +449,7 @@ class DreamFuseInference:
                 scale_factor=scale_factor, device=self.device,
             )
         cond_latent_image_ids = torch.stack(cond_latent_image_ids)
-        print(2)
+
         # Pack condition latents
         cond_image_latents = self._pack_latents(condition_image_latents)
         cond_input = {
@@ -470,7 +462,7 @@ class DreamFuseInference:
         latents, latent_image_ids = self._prepare_latents(
             batch_size, num_channels_latents, height, width, generator
         )
-        print(3)
+
         # Setup timesteps
         sigmas = np.linspace(1.0, 1 / self.config.num_inference_steps, self.config.num_inference_steps)
         image_seq_len = latents.shape[1]
@@ -488,7 +480,7 @@ class DreamFuseInference:
             sigmas=sigmas,
             mu=mu,
         )
-        print(4)
+
         # Denoising loop
         for i, t in enumerate(timesteps):
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
@@ -537,12 +529,12 @@ class DreamFuseInference:
 
             # Compute previous noisy sample
             latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-        print(5)
+
         # Decode latents
         latents = self._unpack_latents(latents, height, width)
         latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
         images = self.vae.decode(latents, return_dict=False)[0]
-        print(6)
+
         # Post-process images
         images = images.add(1).mul(127.5).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy()
         return images
@@ -575,68 +567,3 @@ class DreamFuseInference:
             offset=None
         )
         return latents, latent_image_ids
-
-def main():
-    parser = transformers.HfArgumentParser(InferenceConfig)
-    config: InferenceConfig = parser.parse_args_into_dataclasses()[0]
-    model = DreamFuseInference(config)
-    os.makedirs(config.valid_output_dir, exist_ok=True)
-    for valid_root, valid_json in zip(config.valid_roots, config.valid_jsons):
-        with open(valid_json, 'r') as f:
-            valid_info = json.load(f)
-            
-        # multi gpu
-        to_process = sorted(list(valid_info.keys()))
-
-        # debug
-        to_process = [k for k in to_process if "data_wear" in k and "pixelwave" in k]
-        # debug
-
-        sd_idx = len(to_process) // config.total_num * config.sub_idx
-        ed_idx = len(to_process) // config.total_num * (config.sub_idx + 1)
-        if config.sub_idx < config.total_num - 1:
-            print(config.sub_idx, sd_idx, ed_idx)
-            to_process = to_process[sd_idx:ed_idx]
-        else:
-            print(config.sub_idx, sd_idx)
-            to_process = to_process[sd_idx:]
-        valid_info = {k: valid_info[k] for k in to_process}
-
-        for meta_key, info in tqdm(valid_info.items()):
-            img_name = meta_key.split('/')[-1]
-
-            foreground_img = Image.open(os.path.join(valid_root, info['img_info']['000']))
-            background_img = Image.open(os.path.join(valid_root, info['img_info']['001']))
-            
-            new_fg_mask = Image.open(os.path.join(valid_root, info['img_mask_info']['000_mask_scale']))
-            ori_fg_mask = Image.open(os.path.join(valid_root, info['img_mask_info']['000']))
-
-            # debug
-            foreground_img.save(os.path.join(config.valid_output_dir, f"{img_name}_0.png"))
-            background_img.save(os.path.join(config.valid_output_dir, f"{img_name}_1.png"))
-            ori_fg_mask.save(os.path.join(config.valid_output_dir, f"{img_name}_0_mask.png"))
-            new_fg_mask.save(os.path.join(config.valid_output_dir, f"{img_name}_0_mask_scale.png"))
-            # debug
-
-            foreground_img.paste((255, 255, 255), mask=ImageOps.invert(ori_fg_mask))
-
-            images = model(foreground_img.copy(), background_img.copy(),
-                ori_fg_mask, new_fg_mask,
-                prompt=config.ref_prompts,
-                seed=config.seed,
-                cfg=config.guidance_scale,
-                size_select=config.inference_scale,
-                text_strength=config.text_strength,
-                truecfg=config.truecfg)
-            
-            result_image = Image.fromarray(images[0], "RGB")
-            result_image = result_image.resize(background_img.size)
-            result_image.save(os.path.join(config.valid_output_dir, f"{img_name}_2.png"))
-            # Make grid
-            grid_image = [foreground_img, background_img] + [result_image]
-            result = make_image_grid(grid_image, 1, len(grid_image), size=result_image.size)
-            
-            result.save(os.path.join(config.valid_output_dir, f"{img_name}.jpg"))
-
-if __name__ == "__main__":
-    main()