| import json |
| import os |
| import random |
| import shutil |
| 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 |
| SPACE_LORA_DIR = Path(__file__).resolve().parent / "loras" |
|
|
| _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 _scan_space_loras(): |
| """Return LoRA filenames uploaded to the Space repo under ./loras.""" |
| SPACE_LORA_DIR.mkdir(parents=True, exist_ok=True) |
| return sorted(path.name for path in SPACE_LORA_DIR.glob("*.safetensors") if path.is_file()) |
|
|
|
|
| def _sync_space_loras_to_comfy(lora_dir: Path): |
| """Copy repo LoRAs into ComfyUI's runtime LoRA folder.""" |
| SPACE_LORA_DIR.mkdir(parents=True, exist_ok=True) |
| lora_dir.mkdir(parents=True, exist_ok=True) |
|
|
| for src in SPACE_LORA_DIR.glob("*.safetensors"): |
| if not src.is_file(): |
| continue |
| dst = lora_dir / src.name |
| if (not dst.exists()) or src.stat().st_size != dst.stat().st_size or src.stat().st_mtime > dst.stat().st_mtime: |
| print(f"[lora] syncing {src.name}", flush=True) |
| shutil.copy2(src, dst) |
|
|
|
|
| def _refresh_lora_choices(): |
| """Refresh the Gradio LoRA dropdown from files in ./loras.""" |
| return gr.update(choices=_scan_space_loras(), value=[]) |
|
|
|
|
| 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) |
|
|
| |
| _sync_space_loras_to_comfy(lora_dir) |
|
|
| 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, |
| selected_loras=None, |
| lora_strength=0.8, |
| ): |
| if selected_loras is None: |
| selected_loras = [] |
| elif isinstance(selected_loras, str): |
| selected_loras = [selected_loras] |
|
|
| available = set(_scan_space_loras()) |
| selected_loras = [name for name in selected_loras if name in available] |
|
|
| 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]}, |
| }, |
| } |
|
|
| |
| model_node = ["1", 0] |
| next_node_id = 20 |
| for lora_name in selected_loras: |
| node_id = str(next_node_id) |
| workflow[node_id] = { |
| "class_type": "LoraLoaderModelOnly", |
| "inputs": { |
| "lora_name": lora_name, |
| "strength_model": float(lora_strength), |
| "model": model_node, |
| }, |
| } |
| model_node = [node_id, 0] |
| next_node_id += 1 |
|
|
| workflow["5"]["inputs"]["model"] = model_node |
|
|
| 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, selected_loras=None, lora_strength=0.8): |
| 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=1024, |
| height=1024, |
| steps=10, |
| cfg=1.0, |
| seed=0, |
| randomize_seed=True, |
| sampler="er_sde", |
| scheduler="simple", |
| selected_loras=None, |
| lora_strength=0.8, |
| ): |
| 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, selected_loras, lora_strength) |
| 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=1024, |
| height=1024, |
| 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("# Redcraft Krea2") |
| gr.Markdown("ComfyUI-native Redcraft Krea2 generation and image editing.") |
|
|
| with gr.Tabs(): |
| with gr.Tab("Generate"): |
| with gr.Row(): |
| with gr.Column(scale=5): |
| prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Describe the image to generate.") |
| negative_prompt = gr.Textbox(label="Negative prompt", lines=2, value="") |
| with gr.Row(): |
| lora_selector = gr.Dropdown( |
| choices=_scan_space_loras(), |
| value=[], |
| multiselect=True, |
| label="LoRAs", |
| info="Select one or more .safetensors files uploaded to the loras/ folder.", |
| ) |
| refresh_loras = gr.Button("Refresh LoRAs", scale=0) |
| lora_strength = gr.Slider( |
| 0.0, |
| 1.5, |
| value=0.8, |
| step=0.05, |
| label="LoRA strength", |
| info="Shared model strength for all selected LoRAs.", |
| ) |
| refresh_loras.click(_refresh_lora_choices, outputs=lora_selector) |
| with gr.Row(): |
| width = gr.Slider(512, 1536, value=1024, step=64, label="Width") |
| height = gr.Slider(512, 1536, value=1024, step=64, label="Height") |
| with gr.Row(): |
| steps = gr.Slider(1, 30, value=10, step=1, label="Steps") |
| cfg = gr.Slider(0.0, 8.0, value=1.0, step=0.1, label="CFG") |
| with gr.Row(): |
| seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed") |
| randomize_seed = gr.Checkbox(value=True, label="Randomize seed") |
| with gr.Accordion("Sampler", open=False): |
| sampler = gr.Dropdown( |
| ["er_sde", "euler", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"], |
| value="er_sde", |
| label="Sampler", |
| ) |
| scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="simple", label="Scheduler") |
| run = gr.Button("Generate", variant="primary") |
| with gr.Column(scale=6): |
| output = gr.Image(label="Result", format="png", elem_id="result-image") |
|
|
| inputs = [prompt, negative_prompt, width, height, steps, cfg, seed, randomize_seed, sampler, scheduler, lora_selector, lora_strength] |
| run.click(generate, inputs, [output, seed]) |
| prompt.submit(generate, inputs, [output, seed]) |
|
|
| with gr.Tab("Edit Image"): |
| with gr.Row(): |
| with gr.Column(scale=5): |
| edit_input = gr.Image(type="pil", label="Input image") |
| edit_prompt = gr.Textbox(label="Edit prompt", lines=5, placeholder="Describe the edit while preserving identity.") |
| edit_negative_prompt = gr.Textbox(label="Negative prompt", lines=2, value="") |
| with gr.Row(): |
| edit_width = gr.Slider(512, 1536, value=1024, step=64, label="Width") |
| edit_height = gr.Slider(512, 1536, value=1024, step=64, label="Height") |
| with gr.Row(): |
| edit_steps = gr.Slider(1, 30, value=12, step=1, label="Steps") |
| edit_cfg = gr.Slider(0.0, 8.0, value=1.2, step=0.1, label="CFG") |
| edit_denoise = gr.Slider( |
| 0.05, |
| 0.8, |
| value=0.35, |
| step=0.05, |
| label="Edit strength", |
| info="Lower values preserve identity and composition more strongly.", |
| ) |
| with gr.Row(): |
| edit_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed") |
| edit_randomize_seed = gr.Checkbox(value=True, label="Randomize seed") |
| with gr.Accordion("Sampler", open=False): |
| edit_sampler = gr.Dropdown( |
| ["er_sde", "euler", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"], |
| value="er_sde", |
| label="Sampler", |
| ) |
| edit_scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="simple", label="Scheduler") |
| edit_run = gr.Button("Edit Image", variant="primary") |
| with gr.Column(scale=6): |
| edit_output = gr.Image(label="Edited image", 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("Identity Edit"): |
| gr.Markdown( |
| "Instruction-based, identity-preserving editing using the community LoRA " |
| "[`conradlocke/krea2-identity-edit`](https://huggingface.co/conradlocke/krea2-identity-edit) " |
| "on the Redcraft Krea2 checkpoint, via the " |
| "[ComfyUI-Krea2Edit](https://github.com/lbouaraba/comfyui-krea2edit) node pack. " |
| "Give it an image and a plain-language instruction; it edits while preserving what you " |
| "didn't ask to change, including the person." |
| ) |
| with gr.Row(): |
| with gr.Column(scale=5): |
| id_input = gr.Image(type="pil", label="Source image") |
| id_prompt = gr.Textbox( |
| label="Edit instruction", |
| lines=3, |
| placeholder="e.g. create a photo of this person at a night market", |
| ) |
| id_ref_boost = gr.Slider( |
| 0.0, |
| 10.0, |
| value=4.0, |
| step=0.5, |
| label="Likeness (ref_boost)", |
| info="How hard the edit holds the reference. 1 = off, 4 = strong likeness (recommended), 8+ over-copies.", |
| ) |
| with gr.Accordion("Advanced settings", open=False): |
| id_grounding_px = gr.Slider( |
| 384, |
| 1536, |
| value=768, |
| step=64, |
| label="Grounding resolution (px)", |
| info="Lower = stronger edit adherence. Higher = stronger identity likeness. Trained range 384-768; 1024+ often still works.", |
| ) |
| id_max_mp = gr.Slider(0.5, 2.0, value=1.0, step=0.1, label="Output size (megapixels)") |
| with gr.Row(): |
| id_steps = gr.Slider(4, 28, value=10, step=1, label="Steps") |
| id_cfg = gr.Slider( |
| 0.0, |
| 5.0, |
| value=1.0, |
| step=0.5, |
| label="CFG", |
| info="Turbo convention: 1.0 = guidance off. Raise (e.g. 3) with more steps for removals/large edits.", |
| ) |
| with gr.Row(): |
| id_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed") |
| id_randomize_seed = gr.Checkbox(value=True, label="Randomize seed") |
| with gr.Row(): |
| id_sampler = gr.Dropdown( |
| ["euler", "er_sde", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"], |
| value="euler", |
| label="Sampler", |
| ) |
| id_scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="simple", label="Scheduler") |
| id_run = gr.Button("Edit Identity", variant="primary") |
| with gr.Column(scale=6): |
| id_output = gr.Image(label="Edited image", 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() |
|
|
