Use ComfyUI Krea2 component loaders

app.py

@@ -17,6 +17,9 @@ from PIL import Image
 
 MODEL_REPO = "Daankular/redcraft-krea2-fp8"
 MODEL_FILE = "redcraftKREA2RedMix_krea2Edition.safetensors"
+COMFY_KREA_REPO = "Comfy-Org/Krea-2"
+TEXT_ENCODER_FILE = "qwen3vl_4b_bf16.safetensors"
+VAE_FILE = "qwen_image_vae.safetensors"
 COMFY_REPO = "https://github.com/comfyanonymous/ComfyUI.git"
 COMFY_DIR = Path(os.environ.get("COMFYUI_DIR", "/tmp/ComfyUI"))
 COMFY_HOST = "127.0.0.1"
@@ -51,15 +54,20 @@ def _validate_comfyui():
     response = requests.get(f"{COMFY_URL}/object_info", timeout=30)
     response.raise_for_status()
     object_info = response.json()
-    required_nodes = ["CheckpointLoaderSimple", "CLIPTextEncode", "KSampler", "VAEDecode", "SaveImage"]
+    required_nodes = ["UNETLoader", "CLIPLoader", "VAELoader", "CLIPTextEncode", "KSampler", "VAEDecode", "SaveImage"]
     missing = [node for node in required_nodes if node not in object_info]
     if missing:
         raise RuntimeError(f"ComfyUI is missing required nodes: {', '.join(missing)}")
 
-    checkpoint_info = object_info["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0]
-    if MODEL_FILE not in checkpoint_info:
-        available = ", ".join(checkpoint_info[:10])
-        raise RuntimeError(f"Redcraft checkpoint is not visible to ComfyUI. First checkpoints: {available}")
+    unet_info = object_info["UNETLoader"]["input"]["required"]["unet_name"][0]
+    clip_info = object_info["CLIPLoader"]["input"]["required"]["clip_name"][0]
+    vae_info = object_info["VAELoader"]["input"]["required"]["vae_name"][0]
+    if MODEL_FILE not in unet_info:
+        raise RuntimeError(f"Redcraft diffusion model is not visible to ComfyUI. First models: {', '.join(unet_info[:10])}")
+    if TEXT_ENCODER_FILE not in clip_info:
+        raise RuntimeError(f"Krea2 text encoder is not visible to ComfyUI. First encoders: {', '.join(clip_info[:10])}")
+    if VAE_FILE not in vae_info:
+        raise RuntimeError(f"Krea2 VAE is not visible to ComfyUI. First VAEs: {', '.join(vae_info[:10])}")
 
 
 def _ensure_comfyui():
@@ -83,12 +91,28 @@ def _ensure_comfyui():
             _run([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], cwd=COMFY_DIR)
             marker.write_text("ok", encoding="utf-8")
 
-        checkpoint_dir = COMFY_DIR / "models" / "checkpoints"
-        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        diffusion_dir = COMFY_DIR / "models" / "diffusion_models"
+        text_encoder_dir = COMFY_DIR / "models" / "text_encoders"
+        vae_dir = COMFY_DIR / "models" / "vae"
+        diffusion_dir.mkdir(parents=True, exist_ok=True)
+        text_encoder_dir.mkdir(parents=True, exist_ok=True)
+        vae_dir.mkdir(parents=True, exist_ok=True)
         hf_hub_download(
             repo_id=MODEL_REPO,
             filename=MODEL_FILE,
-            local_dir=str(checkpoint_dir),
+            local_dir=str(diffusion_dir),
+            token=os.environ.get("HF_TOKEN"),
+        )
+        hf_hub_download(
+            repo_id=COMFY_KREA_REPO,
+            filename=f"text_encoders/{TEXT_ENCODER_FILE}",
+            local_dir=str(COMFY_DIR / "models"),
+            token=os.environ.get("HF_TOKEN"),
+        )
+        hf_hub_download(
+            repo_id=COMFY_KREA_REPO,
+            filename=f"vae/{VAE_FILE}",
+            local_dir=str(COMFY_DIR / "models"),
             token=os.environ.get("HF_TOKEN"),
         )
 
@@ -124,12 +148,20 @@ def _start_comfyui():
 def _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sampler, scheduler):
     workflow = {
         "1": {
-            "class_type": "CheckpointLoaderSimple",
-            "inputs": {"ckpt_name": MODEL_FILE},
+            "class_type": "UNETLoader",
+            "inputs": {"unet_name": MODEL_FILE, "weight_dtype": "default"},
+        },
+        "8": {
+            "class_type": "CLIPLoader",
+            "inputs": {"clip_name": TEXT_ENCODER_FILE, "type": "krea2", "device": "default"},
+        },
+        "9": {
+            "class_type": "VAELoader",
+            "inputs": {"vae_name": VAE_FILE},
         },
         "2": {
             "class_type": "CLIPTextEncode",
-            "inputs": {"text": prompt, "clip": ["1", 1]},
+            "inputs": {"text": prompt, "clip": ["8", 0]},
         },
         "4": {
             "class_type": "EmptyLatentImage",
@@ -152,7 +184,7 @@ def _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sa
         },
         "6": {
             "class_type": "VAEDecode",
-            "inputs": {"samples": ["5", 0], "vae": ["1", 2]},
+            "inputs": {"samples": ["5", 0], "vae": ["9", 0]},
         },
         "7": {
             "class_type": "SaveImage",
@@ -163,7 +195,7 @@ def _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sa
     if negative_prompt and negative_prompt.strip():
         workflow["3"] = {
             "class_type": "CLIPTextEncode",
-            "inputs": {"text": negative_prompt, "clip": ["1", 1]},
+            "inputs": {"text": negative_prompt, "clip": ["8", 0]},
         }
     else:
         workflow["3"] = {