import json import os import random import subprocess import sys import threading import time import uuid from pathlib import Path import gradio as gr import requests import spaces from huggingface_hub import hf_hub_download 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" IDENTITY_LORA_REPO = "conradlocke/krea2-identity-edit" IDENTITY_LORA_FILE = "krea2_identity_edit_v1_2.safetensors" KREA2EDIT_NODE_REPO = "https://github.com/lbouaraba/comfyui-krea2edit" KREA2EDIT_NODE_DIRNAME = "comfyui-krea2edit" COMFY_REPO = "https://github.com/comfyanonymous/ComfyUI.git" COMFY_DIR = Path(os.environ.get("COMFYUI_DIR", "/tmp/ComfyUI")) COMFY_HOST = "127.0.0.1" COMFY_PORT = int(os.environ.get("COMFYUI_PORT", "8188")) COMFY_URL = f"http://{COMFY_HOST}:{COMFY_PORT}" MAX_SEED = 2**31 - 1 _comfy_lock = threading.Lock() _comfy_process = None def _run(cmd, cwd=None): print("[setup]", " ".join(map(str, cmd)), flush=True) subprocess.check_call(cmd, cwd=str(cwd) if cwd else None) def _wait_for_comfy(timeout=180): deadline = time.time() + timeout last_error = None while time.time() < deadline: try: response = requests.get(f"{COMFY_URL}/system_stats", timeout=2) if response.ok: return except Exception as exc: last_error = exc time.sleep(1) raise RuntimeError(f"ComfyUI did not start in time: {last_error}") def _validate_comfyui(): response = requests.get(f"{COMFY_URL}/object_info", timeout=30) response.raise_for_status() object_info = response.json() 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)}") 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])}") identity_edit_nodes = ["LoraLoaderModelOnly", "Krea2EditModelPatch", "Krea2EditGroundedEncode", "EmptySD3LatentImage"] missing_identity_nodes = [node for node in identity_edit_nodes if node not in object_info] if missing_identity_nodes: raise RuntimeError( f"ComfyUI-Krea2Edit nodes are missing: {', '.join(missing_identity_nodes)}. " f"Check the {KREA2EDIT_NODE_REPO} custom node install." ) lora_info = object_info["LoraLoaderModelOnly"]["input"]["required"]["lora_name"][0] if IDENTITY_LORA_FILE not in lora_info: raise RuntimeError(f"Identity-edit LoRA is not visible to ComfyUI. First loras: {', '.join(lora_info[:10])}") def _ensure_comfyui(): global _comfy_process with _comfy_lock: if _comfy_process is not None and _comfy_process.poll() is None: return try: response = requests.get(f"{COMFY_URL}/system_stats", timeout=2) if response.ok: return except Exception: pass if not COMFY_DIR.exists(): _run(["git", "clone", "--depth", "1", COMFY_REPO, str(COMFY_DIR)]) marker = COMFY_DIR / ".requirements-installed" if not marker.exists(): _run([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], cwd=COMFY_DIR) marker.write_text("ok", encoding="utf-8") krea2edit_dir = COMFY_DIR / "custom_nodes" / KREA2EDIT_NODE_DIRNAME if not krea2edit_dir.exists(): _run(["git", "clone", "--depth", "1", KREA2EDIT_NODE_REPO, str(krea2edit_dir)]) diffusion_dir = COMFY_DIR / "models" / "diffusion_models" text_encoder_dir = COMFY_DIR / "models" / "text_encoders" vae_dir = COMFY_DIR / "models" / "vae" lora_dir = COMFY_DIR / "models" / "loras" 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) lora_dir.mkdir(parents=True, exist_ok=True) hf_hub_download( repo_id=MODEL_REPO, filename=MODEL_FILE, 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"), ) hf_hub_download( repo_id=IDENTITY_LORA_REPO, filename=IDENTITY_LORA_FILE, local_dir=str(lora_dir), token=os.environ.get("HF_TOKEN"), ) def _start_comfyui(): global _comfy_process with _comfy_lock: if _comfy_process is not None and _comfy_process.poll() is None: return try: response = requests.get(f"{COMFY_URL}/system_stats", timeout=2) if response.ok: _validate_comfyui() return except Exception: pass cmd = [ sys.executable, "main.py", "--listen", COMFY_HOST, "--port", str(COMFY_PORT), "--disable-auto-launch", ] _comfy_process = subprocess.Popen(cmd, cwd=str(COMFY_DIR)) _wait_for_comfy() _validate_comfyui() def _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sampler, scheduler): workflow = { "1": { "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": ["8", 0]}, }, "4": { "class_type": "EmptyLatentImage", "inputs": {"width": int(width), "height": int(height), "batch_size": 1}, }, "5": { "class_type": "KSampler", "inputs": { "seed": int(seed), "steps": int(steps), "cfg": float(cfg), "sampler_name": sampler, "scheduler": scheduler, "denoise": 1.0, "model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["4", 0], }, }, "6": { "class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["9", 0]}, }, "7": { "class_type": "SaveImage", "inputs": {"filename_prefix": "redcraft", "images": ["6", 0]}, }, } if negative_prompt and negative_prompt.strip(): workflow["3"] = { "class_type": "CLIPTextEncode", "inputs": {"text": negative_prompt, "clip": ["8", 0]}, } else: workflow["3"] = { "class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["2", 0]}, } return workflow def _upload_image_to_comfy(image): if image is None: raise ValueError("Upload an image to edit.") image = image.convert("RGB") filename = f"redcraft-input-{uuid.uuid4().hex}.png" temp_path = Path("/tmp") / filename image.save(temp_path) with temp_path.open("rb") as handle: response = requests.post( f"{COMFY_URL}/upload/image", files={"image": (filename, handle, "image/png")}, data={"overwrite": "true"}, timeout=120, ) response.raise_for_status() data = response.json() return data.get("name", filename) def _resize_for_edit(image, width, height): width, height = int(width), int(height) if width <= 0 or height <= 0: return image.convert("RGB") return image.convert("RGB").resize((width, height), Image.LANCZOS) def _target_size_from_source(image, max_megapixels=1.0): multiple = 16 width, height = image.size megapixels = (width * height) / 1_000_000 if megapixels > max_megapixels: scale = (max_megapixels / megapixels) ** 0.5 width, height = round(width * scale), round(height * scale) width = max(multiple, (width // multiple) * multiple) height = max(multiple, (height // multiple) * multiple) return width, height def _build_edit_workflow(input_filename, prompt, negative_prompt, steps, cfg, seed, sampler, scheduler, denoise): workflow = { "1": { "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}, }, "10": { "class_type": "LoadImage", "inputs": {"image": input_filename}, }, "11": { "class_type": "VAEEncode", "inputs": {"pixels": ["10", 0], "vae": ["9", 0]}, }, "2": { "class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["8", 0]}, }, "5": { "class_type": "KSampler", "inputs": { "seed": int(seed), "steps": int(steps), "cfg": float(cfg), "sampler_name": sampler, "scheduler": scheduler, "denoise": float(denoise), "model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["11", 0], }, }, "6": { "class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["9", 0]}, }, "7": { "class_type": "SaveImage", "inputs": {"filename_prefix": "redcraft-edit", "images": ["6", 0]}, }, } if negative_prompt and negative_prompt.strip(): workflow["3"] = { "class_type": "CLIPTextEncode", "inputs": {"text": negative_prompt, "clip": ["8", 0]}, } else: workflow["3"] = { "class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["2", 0]}, } return workflow def _build_identity_edit_workflow(input_filename, prompt, width, height, steps, cfg, seed, ref_boost, grounding_px, sampler, scheduler): return { "1": { "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}, }, "20": { "class_type": "LoraLoaderModelOnly", "inputs": {"lora_name": IDENTITY_LORA_FILE, "strength_model": 1.0, "model": ["1", 0]}, }, "10": { "class_type": "LoadImage", "inputs": {"image": input_filename}, }, "11": { "class_type": "VAEEncode", "inputs": {"pixels": ["10", 0], "vae": ["9", 0]}, }, "21": { "class_type": "Krea2EditModelPatch", "inputs": { "model": ["20", 0], "source_latent": ["11", 0], "ref_boost": float(ref_boost), "fit_mode": "fit", "vae": ["9", 0], "source_image": ["10", 0], }, }, "2": { "class_type": "Krea2EditGroundedEncode", "inputs": {"clip": ["8", 0], "prompt": prompt, "image": ["10", 0], "grounding_px": int(grounding_px)}, }, "3": { "class_type": "Krea2EditGroundedEncode", "inputs": {"clip": ["8", 0], "prompt": "", "image": ["10", 0], "grounding_px": int(grounding_px)}, }, "4": { "class_type": "EmptySD3LatentImage", "inputs": {"width": int(width), "height": int(height), "batch_size": 1}, }, "5": { "class_type": "KSampler", "inputs": { "seed": int(seed), "steps": int(steps), "cfg": float(cfg), "sampler_name": sampler, "scheduler": scheduler, "denoise": 1.0, "model": ["21", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["4", 0], }, }, "6": { "class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["9", 0]}, }, "7": { "class_type": "SaveImage", "inputs": {"filename_prefix": "identity-edit", "images": ["6", 0]}, }, } def _queue_prompt(workflow): payload = {"prompt": workflow, "client_id": str(uuid.uuid4())} response = requests.post(f"{COMFY_URL}/prompt", json=payload, timeout=30) if not response.ok: raise RuntimeError(f"ComfyUI prompt error {response.status_code}: {response.text[:1000]}") return response.json()["prompt_id"] def _wait_for_history(prompt_id, timeout=900): deadline = time.time() + timeout while time.time() < deadline: response = requests.get(f"{COMFY_URL}/history/{prompt_id}", timeout=30) response.raise_for_status() history = response.json() if prompt_id in history: item = history[prompt_id] status = item.get("status", {}) if status.get("completed"): return item messages = status.get("messages") or [] for message in messages: if isinstance(message, list) and message and message[0] == "execution_error": raise RuntimeError(json.dumps(message[1], indent=2)[:2000]) time.sleep(1) raise RuntimeError("Timed out waiting for ComfyUI generation.") def _load_output_image(history_item): outputs = history_item.get("outputs", {}) for output in outputs.values(): for image in output.get("images", []): params = { "filename": image["filename"], "subfolder": image.get("subfolder", ""), "type": image.get("type", "output"), } response = requests.get(f"{COMFY_URL}/view", params=params, timeout=120) response.raise_for_status() temp_path = Path("/tmp") / f"{uuid.uuid4().hex}.png" temp_path.write_bytes(response.content) return Image.open(temp_path).convert("RGB") raise RuntimeError("ComfyUI completed without returning an image.") def _duration(prompt, negative_prompt, width, height, steps, cfg, seed, randomize_seed, sampler, scheduler): megapixels = max(1.0, (int(width) * int(height)) / (1024 * 1024)) return int(900 + int(steps) * 8 * megapixels) def _edit_duration( input_image, prompt, negative_prompt, width, height, steps, cfg, denoise, seed, randomize_seed, sampler, scheduler, ): megapixels = max(1.0, (int(width) * int(height)) / (1024 * 1024)) return int(900 + int(steps) * 9 * megapixels) def _identity_edit_duration( input_image, prompt, ref_boost, grounding_px, max_megapixels, steps, cfg, seed, randomize_seed, sampler, scheduler, ): return int(900 + int(steps) * 10 * max(1.0, float(max_megapixels))) @spaces.GPU(duration=_duration) def generate( prompt, negative_prompt="", width=2048, height=2048, steps=10, cfg=1.0, seed=0, randomize_seed=True, sampler="er_sde", scheduler="simple", ): if not prompt or not prompt.strip(): raise gr.Error("Enter a prompt.") if randomize_seed: seed = random.randint(0, MAX_SEED) try: _start_comfyui() workflow = _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sampler, scheduler) prompt_id = _queue_prompt(workflow) history_item = _wait_for_history(prompt_id) return _load_output_image(history_item), seed except Exception as exc: raise gr.Error(str(exc)) from exc @spaces.GPU(duration=_edit_duration) def edit_image( input_image, prompt, negative_prompt="", width=2048, height=2048, steps=12, cfg=1.2, denoise=0.35, seed=0, randomize_seed=True, sampler="er_sde", scheduler="simple", ): if input_image is None: raise gr.Error("Upload an image to edit.") if not prompt or not prompt.strip(): raise gr.Error("Enter an edit prompt.") if randomize_seed: seed = random.randint(0, MAX_SEED) try: _start_comfyui() resized = _resize_for_edit(input_image, width, height) input_filename = _upload_image_to_comfy(resized) workflow = _build_edit_workflow( input_filename, prompt, negative_prompt, steps, cfg, seed, sampler, scheduler, denoise, ) prompt_id = _queue_prompt(workflow) history_item = _wait_for_history(prompt_id) return _load_output_image(history_item), seed except Exception as exc: raise gr.Error(str(exc)) from exc @spaces.GPU(duration=_identity_edit_duration) def identity_edit( input_image, prompt, ref_boost=4.0, grounding_px=768, max_megapixels=1.0, steps=10, cfg=1.0, seed=0, randomize_seed=True, sampler="euler", scheduler="simple", ): if input_image is None: raise gr.Error("Upload an image to edit.") if not prompt or not prompt.strip(): raise gr.Error("Enter an edit instruction.") if randomize_seed: seed = random.randint(0, MAX_SEED) try: _start_comfyui() source = input_image.convert("RGB") width, height = _target_size_from_source(source, max_megapixels) input_filename = _upload_image_to_comfy(source) workflow = _build_identity_edit_workflow( input_filename, prompt, width, height, steps, cfg, seed, ref_boost, grounding_px, sampler, scheduler, ) prompt_id = _queue_prompt(workflow) history_item = _wait_for_history(prompt_id) return _load_output_image(history_item), seed except Exception as exc: raise gr.Error(str(exc)) from exc CSS = """ .gradio-container { max-width: 1120px !important; margin: 0 auto !important; } #result-image { min-height: 520px; } """ with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo: gr.Markdown(" ") gr.Markdown(" ") with gr.Tabs(): with gr.Tab(" "): with gr.Row(): with gr.Column(scale=5): prompt = gr.Textbox(label=" ", lines=6, placeholder=" ") negative_prompt = gr.Textbox(label=" ", lines=2, value="") with gr.Row(): width = gr.Slider(1024, 2048, value=2048, step=512, label=" ") height = gr.Slider(1024, 2048, value=2048, step=512, label=" ") with gr.Row(): steps = gr.Slider(1, 30, value=10, step=1, label=" ") cfg = gr.Slider(0.0, 8.0, value=1.0, step=0.1, label=" ") with gr.Row(): seed = gr.Slider(0, MAX_SEED, value=0, step=1, label=" ") randomize_seed = gr.Checkbox(value=True, label=" ") with gr.Accordion(" ", open=False): sampler = gr.Dropdown( ["er_sde", "euler", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"], value="euler", label=" ", ) scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="normal", label="Scheduler") run = gr.Button("Art", variant="primary") with gr.Column(scale=6): output = gr.Image(label=" ", format="png", elem_id="result-image") inputs = [prompt, negative_prompt, width, height, steps, cfg, seed, randomize_seed, sampler, scheduler] run.click(generate, inputs, [output, seed]) prompt.submit(generate, inputs, [output, seed]) with gr.Tab(" "): with gr.Row(): with gr.Column(scale=5): edit_input = gr.Image(type="pil", label=" ") edit_prompt = gr.Textbox(label=" ", lines=6, placeholder=" ") edit_negative_prompt = gr.Textbox(label=" ", lines=2, value="") with gr.Row(): edit_width = gr.Slider(1024, 2048, value=2048, step=512, label=" ") edit_height = gr.Slider(1024, 2048, value=2048, step=512, label=" ") with gr.Row(): edit_steps = gr.Slider(1, 30, value=12, step=1, label=" ") edit_cfg = gr.Slider(0.0, 8.0, value=1.2, step=0.1, label=" ") edit_denoise = gr.Slider( 0.05, 0.8, value=0.35, step=0.05, label=" ", info=" ", ) with gr.Row(): edit_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label=" ") edit_randomize_seed = gr.Checkbox(value=True, label=" ") with gr.Accordion(" ", open=False): edit_sampler = gr.Dropdown( ["er_sde", "euler", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"], value="euler", label=" ", ) edit_scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="normal", label=" ") edit_run = gr.Button("Art", variant="primary") with gr.Column(scale=6): edit_output = gr.Image(label=" ", format="png", elem_id="result-image") edit_inputs = [ edit_input, edit_prompt, edit_negative_prompt, edit_width, edit_height, edit_steps, edit_cfg, edit_denoise, edit_seed, edit_randomize_seed, edit_sampler, edit_scheduler, ] edit_run.click(edit_image, edit_inputs, [edit_output, edit_seed]) edit_prompt.submit(edit_image, edit_inputs, [edit_output, edit_seed]) with gr.Tab(" "): gr.Markdown( " " ) with gr.Row(): with gr.Column(scale=5): id_input = gr.Image(type="pil", label=" ") id_prompt = gr.Textbox( label=" ", lines=6, placeholder=" ", ) id_ref_boost = gr.Slider( 0.0, 10.0, value=4.0, step=0.5, label=" ", info=" ", ) with gr.Accordion(" ", open=False): id_grounding_px = gr.Slider( 384, 1536, value=768, step=64, label=" ", info=" ", ) id_max_mp = gr.Slider(0.5, 2.0, value=1.0, step=0.1, label=" ") with gr.Row(): id_steps = gr.Slider(4, 28, value=10, step=1, label=" ") id_cfg = gr.Slider( 0.0, 5.0, value=1.0, step=0.5, label=" ", info=" ", ) with gr.Row(): id_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label=" ") id_randomize_seed = gr.Checkbox(value=True, label=" ") with gr.Row(): id_sampler = gr.Dropdown( ["euler", "er_sde", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"], value="euler", label=" ", ) id_scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="normal", label=" ") id_run = gr.Button("Art", variant="primary") with gr.Column(scale=6): id_output = gr.Image(label=" ", format="png", elem_id="result-image") id_inputs = [ id_input, id_prompt, id_ref_boost, id_grounding_px, id_max_mp, id_steps, id_cfg, id_seed, id_randomize_seed, id_sampler, id_scheduler, ] id_run.click(identity_edit, id_inputs, [id_output, id_seed]) id_prompt.submit(identity_edit, id_inputs, [id_output, id_seed]) gr.Examples( examples=[ ["examples/woman.jpg", "create a photo of this person at a busy night market at night"], ["examples/businessman_suit.jpg", "change the suit jacket to a red leather jacket"], ["examples/man_beach.jpg", "make it a vintage film photo with warm golden-hour light"], ], inputs=[id_input, id_prompt], label="Examples", ) _ensure_comfyui() if __name__ == "__main__": demo.queue().launch()