remove freeu support

2024-04-25 22:21:33 +08:00
parent 911f7224b6
commit 017a3d68fd
10 changed files with 33 additions and 326 deletions
--- a/iopaint/model_manager.py
+++ b/iopaint/model_manager.py
@@ -103,7 +103,6 @@ class ModelManager:
            self.switch_brushnet_method(config)

        self.enable_disable_powerpaint_v2(config)
-        self.enable_disable_freeu(config)
        self.enable_disable_lcm_lora(config)
        return self.model(image, mask, config).astype(np.uint8)

@@ -240,22 +239,6 @@ class ModelManager:
            else:
                logger.info("Disable PowerPaintV2")

-    def enable_disable_freeu(self, config: InpaintRequest):
-        if str(self.model.device) == "mps":
-            return
-
-        if self.available_models[self.name].support_freeu:
-            if config.sd_freeu:
-                freeu_config = config.sd_freeu_config
-                self.model.model.enable_freeu(
-                    s1=freeu_config.s1,
-                    s2=freeu_config.s2,
-                    b1=freeu_config.b1,
-                    b2=freeu_config.b2,
-                )
-            else:
-                self.model.model.disable_freeu()
-
    def enable_disable_lcm_lora(self, config: InpaintRequest):
        if self.available_models[self.name].support_lcm_lora:
            # TODO: change this if load other lora is supported
--- a/iopaint/schema.py
+++ b/iopaint/schema.py
@@ -126,16 +126,6 @@ class ModelInfo(BaseModel):
            ModelType.DIFFUSERS_SD,
        ]

-    @computed_field
-    @property
-    def support_freeu(self) -> bool:
-        return self.model_type in [
-            ModelType.DIFFUSERS_SD,
-            ModelType.DIFFUSERS_SDXL,
-            ModelType.DIFFUSERS_SD_INPAINT,
-            ModelType.DIFFUSERS_SDXL_INPAINT,
-        ] or self.name in [INSTRUCT_PIX2PIX_NAME]
-

 class Choices(str, Enum):
    @classmethod
@@ -225,12 +215,6 @@ class SDSampler(str, Enum):
    lcm = "LCM"


-class FREEUConfig(BaseModel):
-    s1: float = 0.9
-    s2: float = 0.2
-    b1: float = 1.2
-    b2: float = 1.4
-

 class PowerPaintTask(Choices):
    text_guided = "text-guided"
@@ -352,12 +336,6 @@ class InpaintRequest(BaseModel):
    sd_outpainting_softness: float = Field(20.0)
    sd_outpainting_space: float = Field(20.0)

-    sd_freeu: bool = Field(
-        False,
-        description="Enable freeu mode. https://huggingface.co/docs/diffusers/main/en/using-diffusers/freeu",
-    )
-    sd_freeu_config: FREEUConfig = FREEUConfig()
-
    sd_lcm_lora: bool = Field(
        False,
        description="Enable lcm-lora mode. https://huggingface.co/docs/diffusers/main/en/using-diffusers/inference_with_lcm#texttoimage",
@@ -433,9 +411,6 @@ class InpaintRequest(BaseModel):
            if values.sd_lcm_lora:
                logger.info("BrushNet is enabled, set sd_lcm_lora=False")
                values.sd_lcm_lora = False
-            if values.sd_freeu:
-                logger.info("BrushNet is enabled, set sd_freeu=False")
-                values.sd_freeu = False

        if values.enable_controlnet:
            logger.info("ControlNet is enabled, set enable_brushnet=False")
--- a/iopaint/tests/test_brushnet.py
+++ b/iopaint/tests/test_brushnet.py
@@ -10,7 +10,7 @@ import pytest
 import torch

 from iopaint.model_manager import ModelManager
-from iopaint.schema import HDStrategy, SDSampler, FREEUConfig, PowerPaintTask
+from iopaint.schema import HDStrategy, SDSampler, PowerPaintTask

 current_dir = Path(__file__).parent.absolute().resolve()
 save_dir = current_dir / "result"
@@ -32,8 +32,6 @@ def test_runway_brushnet(device, sampler):
        prompt="face of a fox, sitting on a bench",
        sd_steps=sd_steps,
        sd_guidance_scale=7.5,
-        sd_freeu=True,
-        sd_freeu_config=FREEUConfig(),
        enable_brushnet=True,
        brushnet_method=SD_BRUSHNET_CHOICES[0],
    )
--- a/iopaint/tests/test_low_mem.py
+++ b/iopaint/tests/test_low_mem.py
@@ -10,7 +10,7 @@ import pytest
 import torch

 from iopaint.model_manager import ModelManager
-from iopaint.schema import HDStrategy, SDSampler, FREEUConfig
+from iopaint.schema import HDStrategy, SDSampler


@pytest.mark.parametrize("device", ["cuda", "mps"])
@@ -75,35 +75,6 @@ def test_runway_sd_lcm_lora_low_mem(device, sampler):
    )


-@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
-@pytest.mark.parametrize("sampler", [SDSampler.ddim])
-def test_runway_sd_freeu(device, sampler):
-    sd_steps = check_device(device)
-    model = ModelManager(
-        name="runwayml/stable-diffusion-inpainting",
-        device=torch.device(device),
-        disable_nsfw=True,
-        sd_cpu_textencoder=False,
-        low_mem=True,
-    )
-    cfg = get_config(
-        strategy=HDStrategy.ORIGINAL,
-        prompt="face of a fox, sitting on a bench",
-        sd_steps=sd_steps,
-        sd_guidance_scale=7.5,
-        sd_freeu=True,
-        sd_freeu_config=FREEUConfig(),
-    )
-    cfg.sd_sampler = sampler
-
-    assert_equal(
-        model,
-        cfg,
-        f"runway_sd_1_5_freeu_device_{device}_low_mem.png",
-        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
-        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
-    )
-

@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
--- a/iopaint/tests/test_sd_model.py
+++ b/iopaint/tests/test_sd_model.py
@@ -11,7 +11,7 @@ import pytest
 import torch

 from iopaint.model_manager import ModelManager
-from iopaint.schema import HDStrategy, SDSampler, FREEUConfig
+from iopaint.schema import HDStrategy, SDSampler

 current_dir = Path(__file__).parent.absolute().resolve()
 save_dir = current_dir / "result"
@@ -90,35 +90,6 @@ def test_runway_sd_lcm_lora(device, sampler):
    )


-@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
-@pytest.mark.parametrize("sampler", [SDSampler.ddim])
-def test_runway_sd_freeu(device, sampler):
-    sd_steps = check_device(device)
-    model = ModelManager(
-        name="runwayml/stable-diffusion-inpainting",
-        device=torch.device(device),
-        disable_nsfw=True,
-        sd_cpu_textencoder=False,
-    )
-    cfg = get_config(
-        strategy=HDStrategy.ORIGINAL,
-        prompt="face of a fox, sitting on a bench",
-        sd_steps=sd_steps,
-        sd_guidance_scale=7.5,
-        sd_freeu=True,
-        sd_freeu_config=FREEUConfig(),
-    )
-    cfg.sd_sampler = sampler
-
-    assert_equal(
-        model,
-        cfg,
-        f"runway_sd_1_5_freeu_device_{device}.png",
-        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
-        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
-    )
-
-
@pytest.mark.parametrize("device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
--- a/iopaint/tests/test_sdxl.py
+++ b/iopaint/tests/test_sdxl.py
@@ -8,7 +8,7 @@ import pytest
 import torch

 from iopaint.model_manager import ModelManager
-from iopaint.schema import HDStrategy, SDSampler, FREEUConfig
+from iopaint.schema import HDStrategy, SDSampler
 from iopaint.tests.test_model import get_config, assert_equal


@@ -76,60 +76,6 @@ def test_sdxl_cpu_text_encoder(device, strategy, sampler):
    )


-@pytest.mark.parametrize("device", ["cuda", "mps"])
-@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
-@pytest.mark.parametrize("sampler", [SDSampler.ddim])
-def test_sdxl_lcm_lora_and_freeu(device, strategy, sampler):
-    sd_steps = check_device(device)
-
-    model = ModelManager(
-        name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
-        device=torch.device(device),
-        disable_nsfw=True,
-        sd_cpu_textencoder=False,
-    )
-    cfg = get_config(
-        strategy=strategy,
-        prompt="face of a fox, sitting on a bench",
-        sd_steps=sd_steps,
-        sd_strength=1.0,
-        sd_guidance_scale=2.0,
-        sd_lcm_lora=True,
-    )
-    cfg.sd_sampler = sampler
-
-    name = f"device_{device}_{sampler}"
-
-    assert_equal(
-        model,
-        cfg,
-        f"sdxl_{name}_lcm_lora.png",
-        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
-        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
-        fx=2,
-        fy=2,
-    )
-
-    cfg = get_config(
-        strategy=strategy,
-        prompt="face of a fox, sitting on a bench",
-        sd_steps=sd_steps,
-        sd_guidance_scale=7.5,
-        sd_freeu=True,
-        sd_freeu_config=FREEUConfig(),
-    )
-
-    assert_equal(
-        model,
-        cfg,
-        f"sdxl_{name}_freeu_device_{device}.png",
-        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
-        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
-        fx=2,
-        fy=2,
-    )
-
-
@pytest.mark.parametrize("device", ["cuda", "mps"])
@pytest.mark.parametrize(
    "rect",