Update injection_main.py

injection_main.py

@@ -473,10 +473,7 @@ def style_image_with_inversion(
 
 if __name__ == "__main__":
 
-    # Load a pipeline
-    pipe = StableDiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-2-1-base"
-    ).to(device)
+    
 
     # pipe = DiffusionPipeline.from_pretrained(
     #     # "playgroundai/playground-v2-1024px-aesthetic",
@@ -487,9 +484,6 @@ if __name__ == "__main__":
     # )
     # pipe.to("cuda")
 
-    # Set up a DDIM scheduler
-    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-
     parser = argparse.ArgumentParser(description="Stable Diffusion with OmegaConf")
     parser.add_argument(
         "--config", type=str, default="config.yaml", help="Path to the config file"
@@ -517,113 +511,19 @@ if __name__ == "__main__":
     # mode = "dataset"
     out_name = ["content_delegation", "style_delegation", "style_out"]
 
-    if mode == "dataset":
-        cfg = OmegaConf.load(config_dir)
-
-        base_output_path = cfg.out_path
-        if not os.path.exists(cfg.out_path):
-            os.makedirs(cfg.out_path)
-        base_output_path = os.path.join(base_output_path, cfg.exp_name)
-
-        experiment_output_path = utils.exp_utils.make_unique_experiment_path(
-            base_output_path
-        )
-
-        # Save the experiment configuration
-        config_file_path = os.path.join(experiment_output_path, "config.yaml")
-        omegaconf.OmegaConf.save(cfg, config_file_path)
-
-        # Seed all
-
-        annotation = json.load(open(cfg.annotation))
-        with open(os.path.join(experiment_output_path, "annotation.json"), "w") as f:
-            json.dump(annotation, f)
-        for i, entry in enumerate(annotation):
-            utils.exp_utils.seed_all(cfg.seed)
-            image_path = entry["image_path"]
-            src_prompt = entry["source_prompt"]
-            tgt_prompt = entry["target_prompt"]
-            resolution = 512 if isinstance(pipe, StableDiffusionXLPipeline) else 512
-            input_image = utils.exp_utils.get_processed_image(
-                image_path, device, resolution
-            )
-
-            prompt_in = [
-                src_prompt,  # reconstruction
-                tgt_prompt,  # uncontrolled style
-                "",  # controlled style
-            ]
-
-            imgs = style_image_with_inversion(
-                pipe,
-                input_image,
-                src_prompt,
-                style_prompt=prompt_in,
-                num_steps=cfg.num_steps,
-                start_step=cfg.start_step,
-                guidance_scale=cfg.style_cfg_scale,
-                disentangle=cfg.disentangle,
-                resnet_mode=cfg.resnet_mode,
-                share_attn=cfg.share_attn,
-                share_cross_attn=cfg.share_cross_attn,
-                share_resnet_layers=cfg.share_resnet_layers,
-                share_attn_layers=cfg.share_attn_layers,
-                share_key=cfg.share_key,
-                share_query=cfg.share_query,
-                share_value=cfg.share_value,
-                use_content_anchor=cfg.use_content_anchor,
-                use_adain=cfg.use_adain,
-                output_dir=experiment_output_path,
-            )
-
-            for j, img in enumerate(imgs):
-                img.save(f"{experiment_output_path}/out_{i}_{out_name[j]}.png")
-                print(
-                    f"Image saved as {experiment_output_path}/out_{i}_{out_name[j]}.png"
-                )
-    elif mode == "cli":
-        cfg = OmegaConf.load(config_dir)
-        utils.exp_utils.seed_all(cfg.seed)
-        image = utils.exp_utils.get_processed_image(args.image_dir, device, 512)
-        tgt_prompt = args.prompt
-        src_prompt = ""
-        prompt_in = [
-            "",  # reconstruction
-            tgt_prompt,  # uncontrolled style
-            "",  # controlled style
-        ]
-        out_dir = "./out"
-        os.makedirs(out_dir, exist_ok=True)
-        imgs = style_image_with_inversion(
-            pipe,
-            image,
-            src_prompt,
-            style_prompt=prompt_in,
-            num_steps=cfg.num_steps,
-            start_step=cfg.start_step,
-            guidance_scale=cfg.style_cfg_scale,
-            disentangle=cfg.disentangle,
-            resnet_mode=cfg.resnet_mode,
-            share_attn=cfg.share_attn,
-            share_cross_attn=cfg.share_cross_attn,
-            share_resnet_layers=cfg.share_resnet_layers,
-            share_attn_layers=cfg.share_attn_layers,
-            share_key=cfg.share_key,
-            share_query=cfg.share_query,
-            share_value=cfg.share_value,
-            use_content_anchor=cfg.use_content_anchor,
-            use_adain=cfg.use_adain,
-            output_dir=out_dir,
-        )
-        image_base_name = os.path.basename(args.image_dir).split(".")[0]
-        for j, img in enumerate(imgs):
-            img.save(f"{out_dir}/{image_base_name}_out_{out_name[j]}.png")
-            print(f"Image saved as {out_dir}/{image_base_name}_out_{out_name[j]}.png")
-    elif mode == "app":
+    if mode == "app":
         # gradio
         import gradio as gr
         import spaces
 
+        # Load a pipeline
+        pipe = StableDiffusionPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-2-1-base"
+        ).to(device)
+
+        # Set up a DDIM scheduler
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+
         @spaces.GPU()
         def style_transfer_app(
             prompt,