diff --git a/ai_diffusion/client.py b/ai_diffusion/client.py index 4e5cd41fb..8059786c6 100644 --- a/ai_diffusion/client.py +++ b/ai_diffusion/client.py @@ -257,7 +257,7 @@ def find(self, id: ResourceId): if result := self.resources.get(id.string): return result # Fallback to epsilon model if v-prediction model not found - if id.arch is Arch.illu_v: + if id.arch in (Arch.illu_v, Arch.illu_rf): if result := self.resources.get(id._replace(arch=Arch.illu).string): return result # Search for architecture-agnostic model diff --git a/ai_diffusion/control.py b/ai_diffusion/control.py index a283bd8bb..0aee1f003 100644 --- a/ai_diffusion/control.py +++ b/ai_diffusion/control.py @@ -146,7 +146,7 @@ def _update_is_supported(self): if client := root.connection.client_if_connected: models = client.models.for_arch(self._model.arch) - if self.mode.is_ip_adapter and models.arch in [Arch.illu, Arch.illu_v]: + if self.mode.is_ip_adapter and models.arch in [Arch.illu, Arch.illu_v, Arch.illu_rf]: resid = resource_id(ResourceKind.clip_vision, Arch.illu, "ip_adapter") has_clip_vision = client.models.resources.get(resid, None) is not None if not has_clip_vision: @@ -177,7 +177,9 @@ def _update_is_supported(self): model = models.find_control(self.mode) self.has_range = model == models.control.find(self.mode, True) if model is None: - search_arch = Arch.illu if models.arch is Arch.illu_v else models.arch + search_arch = ( + Arch.illu if models.arch in (Arch.illu_v, Arch.illu_rf) else models.arch + ) search_path = ( resources.search_path(ResourceKind.controlnet, search_arch, self.mode) or resources.search_path(ResourceKind.model_patch, search_arch, self.mode) diff --git a/ai_diffusion/resources.py b/ai_diffusion/resources.py index 39af189cc..7745e26d9 100644 --- a/ai_diffusion/resources.py +++ b/ai_diffusion/resources.py @@ -90,6 +90,7 @@ class Arch(Enum): flux2_9b = "Flux 2 Klein 9B" illu = "Illustrious" illu_v = "Illustrious v-prediction" + illu_rf = "Illustrious Rectified Flow" chroma = "Chroma" qwen = "Qwen" qwen_e = "Qwen Edit" @@ -103,6 +104,8 @@ class Arch(Enum): @staticmethod def from_string(string: str, model_type: str = "eps", filename: str | None = None): filename = filename.lower() if filename else "" + if filename == "chenkinnoobxlv02_v02.safetensors": + return Arch.illu_rf if string == "sd15": return Arch.sd15 if string == "sdxl" and model_type == "v-prediction": @@ -123,6 +126,8 @@ def from_string(string: str, model_type: str = "eps", filename: str | None = Non return Arch.illu if string == "illu_v": return Arch.illu_v + if string == "illu_rf": + return Arch.illu_rf if string == "chroma": return Arch.chroma if string == "qwen-image" and "edit" in filename: @@ -175,7 +180,7 @@ def has_controlnet_inpaint(self): @property def supports_regions(self): - return self in [Arch.sd15, Arch.sdxl, Arch.illu, Arch.illu_v] + return self in [Arch.sd15, Arch.sdxl, Arch.illu, Arch.illu_v, Arch.illu_rf] @property def supports_lcm(self): @@ -183,11 +188,11 @@ def supports_lcm(self): @property def supports_clip_skip(self): - return self in [Arch.sd15, Arch.sdxl, Arch.illu, Arch.illu_v] + return self in [Arch.sd15, Arch.sdxl, Arch.illu, Arch.illu_v, Arch.illu_rf] @property def supports_attention_guidance(self): - return self in [Arch.sd15, Arch.sdxl, Arch.illu, Arch.illu_v] + return self in [Arch.sd15, Arch.sdxl, Arch.illu, Arch.illu_v, Arch.illu_rf] @property def supports_cfg(self): @@ -204,7 +209,7 @@ def supports_edit(self): # includes text-to-image models that can also edit @property def is_sdxl_like(self): # illustrious technically uses sdxl architecture, but has a separate ecosystem - return self in [Arch.sdxl, Arch.illu, Arch.illu_v] + return self in [Arch.sdxl, Arch.illu, Arch.illu_v, Arch.illu_rf] @property def is_flux_like(self): @@ -223,7 +228,7 @@ def text_encoders(self): match self: case Arch.sd15: return ["clip_l"] - case Arch.sdxl | Arch.illu | Arch.illu_v: + case Arch.sdxl | Arch.illu | Arch.illu_v | Arch.illu_rf: return ["clip_l", "clip_g"] case Arch.sd3: return ["clip_l", "clip_g"] @@ -257,6 +262,7 @@ def list(): Arch.flux2_9b, Arch.illu, Arch.illu_v, + Arch.illu_rf, Arch.chroma, Arch.qwen, Arch.qwen_e, @@ -721,6 +727,7 @@ def is_required(kind: ResourceKind, arch: Arch, identifier: ControlMode | Upscal resource_id(ResourceKind.controlnet, Arch.sdxl, ControlMode.universal): ["union-sdxl", "xinsirunion"], resource_id(ResourceKind.controlnet, Arch.illu, ControlMode.universal): ["union-sdxl", "xinsirunion"], resource_id(ResourceKind.controlnet, Arch.illu_v, ControlMode.universal): ["union-sdxl", "xinsirunion"], + resource_id(ResourceKind.controlnet, Arch.illu_rf, ControlMode.universal): ["union-sdxl", "xinsirunion"], resource_id(ResourceKind.controlnet, Arch.flux, ControlMode.universal): ["flux.1-dev-controlnet-union-pro-2.0", "flux.1-dev-controlnet-union-pro", "flux.1-dev-controlnet-union", "flux1devcontrolnetunion"], resource_id(ResourceKind.controlnet, Arch.qwen, ControlMode.universal): ["qwen-image-instantx-controlnet-union"], resource_id(ResourceKind.controlnet, Arch.sd15, ControlMode.scribble): ["control_v11p_sd15_scribble", "control_lora_rank128_v11p_sd15_scribble"], @@ -798,6 +805,7 @@ def is_required(kind: ResourceKind, arch: Arch, identifier: ControlMode | Upscal resource_id(ResourceKind.vae, Arch.sdxl, "default"): ["sdxl_vae"], resource_id(ResourceKind.vae, Arch.illu, "default"): ["sdxl_vae"], resource_id(ResourceKind.vae, Arch.illu_v, "default"): ["sdxl_vae"], + resource_id(ResourceKind.vae, Arch.illu_rf, "default"): ["sdxl_vae"], resource_id(ResourceKind.vae, Arch.sd3, "default"): ["sd3"], resource_id(ResourceKind.vae, Arch.flux, "default"): ["flux-", "flux_", "flux/", "flux1", "ae.s"], resource_id(ResourceKind.vae, Arch.flux_k, "default"): ["flux-", "flux_", "flux/", "flux1", "ae.s"], diff --git a/ai_diffusion/ui/server.py b/ai_diffusion/ui/server.py index 118abec4c..c54cf0299 100644 --- a/ai_diffusion/ui/server.py +++ b/ai_diffusion/ui/server.py @@ -371,7 +371,7 @@ def __init__(self, server: Server, parent=None): ), "illu": PackageGroupWidget( _("Illustrious/NoobAI XL models"), - [m for m in optional_models if m.arch in [Arch.illu, Arch.illu_v]], + [m for m in optional_models if m.arch in [Arch.illu, Arch.illu_v, Arch.illu_rf]], is_checkable=True, is_expanded=False, parent=self, diff --git a/ai_diffusion/ui/settings.py b/ai_diffusion/ui/settings.py index a73d80f8e..8a8556319 100644 --- a/ai_diffusion/ui/settings.py +++ b/ai_diffusion/ui/settings.py @@ -644,7 +644,7 @@ def model_name(id: ResourceId, with_file=False): text += "" text += _("Detected base models:") + "\n