Spaces:
Running on Zero
Running on Zero
| # Krea 2 Eric Insert: photo + request -> the same photo with er1cshen added. | |
| # The pipeline is workflow_api.json (open it in ComfyUI to see it), run by | |
| # ComfyUI's own executor. This file only feeds it: the photo, the request, the | |
| # seed, the checked placement phrase; plus resizing to 2 MP and an NSFW check | |
| # (Krea 2 license §4.2). Models load once at startup, as ZeroGPU recommends: | |
| # https://huggingface.co/blog/run-comfyui-workflows-on-spaces | |
| import copy | |
| import json | |
| import os | |
| import re | |
| import sys | |
| import tempfile | |
| import uuid | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image | |
| from transformers import pipeline | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| COMFYUI = os.path.join(HERE, "ComfyUI") | |
| MODELS = os.path.join(COMFYUI, "models") | |
| PIXELS = 2 * 1024 * 1024 # every photo is resized to this; smaller faces come out with artifacts | |
| DEFAULT_PLACE = "standing next to the person" | |
| EMPTY_PLACE = "standing in the scene" # used when the photo has no people, so the model doesn't invent one to stand next to | |
| PERSON_WORDS = re.compile(r"\b(him|her|them|his|their|person|people|next to|beside)\b") | |
| STARTS = {"on", "in", "at", "behind", "beside", "next", "near", "between", "under", "above", "by", "with", "from", "inside", "outside", "across", "against", "along", "around", "among", "to"} | |
| CHANGE_WORDS = re.compile(r"\b(turn|turning|transform|transforming|into|remove|removing|replace|replacing|change|changing|make|making)\b") | |
| for f in ( | |
| "diffusion_models/krea2_turbo_bf16.safetensors", | |
| "text_encoders/qwen3vl_4b_bf16.safetensors", | |
| "vae/qwen_image_vae.safetensors", | |
| ): | |
| hf_hub_download("Comfy-Org/Krea-2", f, local_dir=MODELS) | |
| hf_hub_download( | |
| "yuncuns/krea2-eric-insert", | |
| "krea2-eric-insert.safetensors", | |
| local_dir=os.path.join(MODELS, "loras"), | |
| ) | |
| sys.path.insert(0, COMFYUI) | |
| import asyncio # noqa: E402 | |
| import execution # noqa: E402 | |
| import folder_paths # noqa: E402 | |
| import nodes # noqa: E402 | |
| import server # noqa: E402 | |
| from comfy import model_management # noqa: E402 | |
| loop = asyncio.new_event_loop() | |
| asyncio.set_event_loop(loop) | |
| server_instance = server.PromptServer(loop) | |
| execution.PromptQueue(server_instance) | |
| asyncio.run(nodes.init_extra_nodes()) | |
| # Same arguments ComfyUI's main.py passes for --cache-classic. | |
| executor = execution.PromptExecutor( | |
| server_instance, cache_type=execution.CacheType.CLASSIC, cache_args={"lru": 0, "ram": 0, "ram_inactive": 0} | |
| ) | |
| WORKFLOW = json.load(open(os.path.join(HERE, "workflow_api.json"))) | |
| def node(title: str) -> str: | |
| return next(k for k, v in WORKFLOW.items() if v["_meta"]["title"] == title) | |
| PHOTO, REQUEST, REWRITE_OUT, PLACEMENT, PROMPT_START, PROMPT_END, SAMPLER, OUTPUT = map( | |
| node, ("photo", "request", "rewrite output", "placement", "prompt start", "prompt end", "sampler", "output") | |
| ) | |
| def run(prompt: dict, outputs: list[str]) -> dict: | |
| executor.execute(prompt, str(uuid.uuid4()), {}, outputs) | |
| if not executor.success: | |
| raise gr.Error(f"Generation failed: {executor.status_messages[-1:]}") | |
| return executor.history_result["outputs"] | |
| class GPUPreload: | |
| """Startup-only output node: loads the graph's models onto the GPU once.""" | |
| def INPUT_TYPES(cls): | |
| return {"required": {"model": ("MODEL",), "clip": ("CLIP",), "vae": ("VAE",)}} | |
| RETURN_TYPES = () | |
| FUNCTION = "load" | |
| OUTPUT_NODE = True | |
| def load(self, model, clip, vae): | |
| model_management.load_models_gpu([model, clip.patcher, vae.patcher]) | |
| return {} | |
| nodes.NODE_CLASS_MAPPINGS["GPUPreload"] = GPUPreload | |
| patch_inputs = WORKFLOW[node("Krea2 Edit (source patch)")]["inputs"] | |
| encode_inputs = WORKFLOW[node("Krea2 Edit (grounded encode)")]["inputs"] | |
| warmup = copy.deepcopy(WORKFLOW) | |
| warmup["preload"] = { | |
| "class_type": "GPUPreload", | |
| "inputs": {"model": patch_inputs["model"], "clip": encode_inputs["clip"], "vae": patch_inputs["vae"]}, | |
| "_meta": {"title": "preload"}, | |
| } | |
| # The executor always runs under torch.inference_mode(); weights loaded that way | |
| # become inference tensors that ZeroGPU cannot hand to its GPU worker | |
| # ("Inference tensors do not track version counter"). Load them with it off. | |
| _inference_mode = torch.inference_mode | |
| torch.inference_mode = lambda *args, **kwargs: _inference_mode(False) | |
| try: | |
| run(warmup, ["preload"]) | |
| finally: | |
| torch.inference_mode = _inference_mode | |
| nsfw = pipeline("image-classification", model="Falconsai/nsfw_image_detection") | |
| people = pipeline("object-detection", model="hustvl/yolos-tiny", device="cpu") | |
| def check_placement(text: str) -> str: | |
| phrase = text.split("</think>")[-1].strip().strip(".\"'").lower() | |
| first = phrase.split(" ", 1)[0] | |
| if ( | |
| not re.fullmatch(r"[a-z][a-z ,'-]{2,80}", phrase) | |
| or not (first.endswith("ing") or first in STARTS) | |
| or CHANGE_WORDS.search(phrase) | |
| ): | |
| return DEFAULT_PLACE | |
| return phrase | |
| def generate(image: Image.Image, request: str, seed: int) -> tuple[Image.Image, str]: | |
| if image is None: | |
| raise gr.Error("Add a photo.") | |
| image = image.convert("RGB") | |
| # sides must be multiples of 16, or the model leaves a strip of noise along the edge | |
| scale = (PIXELS / (image.width * image.height)) ** 0.5 | |
| image = image.resize((int(image.width * scale) // 16 * 16, int(image.height * scale) // 16 * 16), Image.LANCZOS) | |
| with tempfile.NamedTemporaryFile( | |
| suffix=".png", dir=folder_paths.get_input_directory(), delete=False | |
| ) as f: | |
| image.save(f.name) | |
| path = f.name | |
| prompt = copy.deepcopy(WORKFLOW) | |
| prompt[REQUEST]["inputs"]["value"] = request.strip()[:300] | |
| text = run(prompt, [REWRITE_OUT])[REWRITE_OUT]["text"][0] | |
| placement = check_placement(text) | |
| if PERSON_WORDS.search(placement) and not any(d["label"] == "person" for d in people(image, threshold=0.7)): | |
| placement = EMPTY_PLACE | |
| prompt[PLACEMENT]["inputs"]["value"] = placement | |
| # the prompt the edit model gets, shown under the result and in the logs (the workflow joins the same three parts) | |
| final = prompt[PROMPT_START]["inputs"]["value"] + placement + prompt[PROMPT_END]["inputs"]["value"] | |
| print(f"seed {seed} · request: {request.strip()[:300]}\nprompt: {final}", flush=True) | |
| prompt[PHOTO]["inputs"]["image"] = os.path.basename(path) | |
| prompt[SAMPLER]["inputs"]["seed"] = int(seed) | |
| result = run(prompt, [OUTPUT])[OUTPUT]["images"][0] | |
| os.unlink(path) | |
| out_path = os.path.join(folder_paths.get_temp_directory(), result["subfolder"], result["filename"]) | |
| out = Image.open(out_path).convert("RGB") | |
| os.unlink(out_path) | |
| if any(r["label"] == "nsfw" and r["score"] > 0.5 for r in nsfw(out)): | |
| raise gr.Error("Blocked by the content filter.") | |
| return out, final | |
| REQUEST_HINT = WORKFLOW[REQUEST]["inputs"]["value"] | |
| SEED = WORKFLOW[SAMPLER]["inputs"]["seed"] | |
| demo = gr.Interface( | |
| fn=generate, | |
| inputs=[ | |
| gr.Image(type="pil", label="Photo"), | |
| gr.Textbox(value=REQUEST_HINT, label="Prompt", lines=2), | |
| gr.Number(value=SEED, label="Seed", precision=0), | |
| ], | |
| outputs=[gr.Image(label="Result", format="png"), gr.Textbox(label="Final prompt", interactive=False)], | |
| examples=[ | |
| ["reagan.png", REQUEST_HINT, SEED], | |
| ["einstein.jpg", REQUEST_HINT, SEED], | |
| ["obama-biden.jpg", "insert a ridiculously handsome man walking with them", SEED], | |
| ["apollo.jpg", "insert a ridiculously handsome man", SEED], | |
| ], | |
| cache_examples=False, | |
| flagging_mode="never", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |