app.py

import base64
import io
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
import piexif

MODEL_REPO = "beznogim666/test2"
MODEL_FILE = "model.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"
COMFY_PORT = int(os.environ.get("COMFYUI_PORT", "8188"))
COMFY_URL = f"http://{COMFY_HOST}:{COMFY_PORT}"
MAX_SEED = 2**31 - 1

LORA_REPO = "beznogim666/test2"
LORA_FILE = "lora1.safetensors"
LORA_NONE = "none"

DURATION_MIN = 10
DURATION_MAX = 300
DURATION_DEFAULT = 90

_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",
        "LoraLoaderModelOnly",
    ]
    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])}")

    lora_info = object_info["LoraLoaderModelOnly"]["input"]["required"]["lora_name"][0]
    if LORA_FILE not in lora_info:
        raise RuntimeError(f"LoRA is not visible to ComfyUI. First LoRAs: {', '.join(lora_info[:10])}")


def _ensure_comfyui():
    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")

    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=LORA_REPO,
        filename=LORA_FILE,
        local_dir=str(lora_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"),
    )


def _start_comfyui():
    global _comfy_process

    with _comfy_lock:
        if _comfy_process is not None and _comfy_process.poll() is None:
            try:
                response = requests.get(f"{COMFY_URL}/system_stats", timeout=2)
                if response.ok:
                    return
            except Exception:
                pass

        cmd = [
            sys.executable,
            "main.py",
            "--listen",
            COMFY_HOST,
            "--port",
            str(COMFY_PORT),
            "--disable-auto-launch",
            "--use-sage-attention",
        ]
        _comfy_process = subprocess.Popen(cmd, cwd=str(COMFY_DIR))
        _wait_for_comfy()
        _validate_comfyui()


def _list_loras():
    try:
        response = requests.get(f"{COMFY_URL}/object_info/LoraLoaderModelOnly", timeout=5)
        if response.ok:
            info = response.json()
            names = info["LoraLoaderModelOnly"]["input"]["required"]["lora_name"][0]
            return [LORA_NONE] + list(names)
    except Exception:
        pass

    lora_dir = COMFY_DIR / "models" / "loras"
    if lora_dir.exists():
        names = sorted(p.name for p in lora_dir.glob("*.safetensors"))
        return [LORA_NONE] + names
    return [LORA_NONE]


def _refresh_comfy_models():
    """
    Форсит обновление внутреннего кэша путей ComfyUI, 
    чтобы свежескачанные модели/лоры сразу стали доступны API-лоадерам.
    """
    try:
        response = requests.post(f"{COMFY_URL}/extra_model_paths", json={}, timeout=5)
        if response.ok:
            print("[setup] ComfyUI model cache refreshed successfully.", flush=True)
    except Exception as e:
        print(f"[error] Failed to refresh ComfyUI model cache: {e}", flush=True)


def _cleanup_old_files(directory, keep_filenames):
    """
    Удаляет все файлы в указанной директории, кроме базовых дефолтных
    и того кастомного файла, который мы используем прямо сейчас.
    """
    try:
        directory = Path(directory)
        if not directory.exists():
            return
            
        keep_set = {f.lower() for f in keep_filenames if f}
        
        for file_path in directory.glob("*"):
            if file_path.is_file():
                # Не трогаем скрытые/системные файлы (.gitattributes и т.д.)
                if file_path.name.startswith("."):
                    continue
                    
                if file_path.name.lower() not in keep_set:
                    try:
                        file_path.unlink()
                        print(f"[ZDR] Cleaned up old file from disk: {file_path}", flush=True)
                    except Exception as e:
                        print(f"[ZDR] Failed to delete old file {file_path}: {e}", flush=True)
    except Exception as e:
        print(f"[error] Error during disk cleanup: {e}", flush=True)


def _download_dynamic_model(repo, filename):
    if not repo or not filename:
        # Если кастомное поле пустое — чистим старые скачанные модели, оставляя только дефолт
        _cleanup_old_files(COMFY_DIR / "models" / "diffusion_models", [MODEL_FILE])
        return MODEL_FILE
    repo = repo.strip()
    filename = filename.strip()
    if not repo or not filename:
        _cleanup_old_files(COMFY_DIR / "models" / "diffusion_models", [MODEL_FILE])
        return MODEL_FILE

    diffusion_dir = COMFY_DIR / "models" / "diffusion_models"
    diffusion_dir.mkdir(parents=True, exist_ok=True)

    try:
        print(f"[setup] Downloading dynamic Checkpoint: {repo}/{filename}", flush=True)
        local_path = hf_hub_download(
            repo_id=repo,
            filename=filename,
            local_dir=str(diffusion_dir),
            token=os.environ.get("HF_TOKEN"),
        )
        downloaded_name = Path(local_path).name
        
        # Оставляем на диске только базовую модель и ту, что только что скачали
        _cleanup_old_files(diffusion_dir, [MODEL_FILE, downloaded_name])
        
        _refresh_comfy_models()
        return str(Path(local_path).relative_to(diffusion_dir))
    except Exception as e:
        print(f"[error] Failed to download dynamic Checkpoint from {repo}/{filename}: {e}", flush=True)
        raise RuntimeError(f"Failed to download dynamic Checkpoint: {e}")


def _download_dynamic_lora(repo, filename):
    if not repo or not filename:
        # Чистим старые динамические лоры, оставляя только дефолт
        _cleanup_old_files(COMFY_DIR / "models" / "loras", [LORA_FILE])
        return None
    repo = repo.strip()
    filename = filename.strip()
    if not repo or not filename:
        _cleanup_old_files(COMFY_DIR / "models" / "loras", [LORA_FILE])
        return None

    lora_dir = COMFY_DIR / "models" / "loras"
    lora_dir.mkdir(parents=True, exist_ok=True)

    try:
        print(f"[setup] Downloading dynamic LoRA: {repo}/{filename}", flush=True)
        local_path = hf_hub_download(
            repo_id=repo,
            filename=filename,
            local_dir=str(lora_dir),
            token=os.environ.get("HF_TOKEN"),
        )
        downloaded_name = Path(local_path).name
        
        # Оставляем только базовую лору и новую скачанную
        _cleanup_old_files(lora_dir, [LORA_FILE, downloaded_name])
        
        _refresh_comfy_models()
        return str(Path(local_path).relative_to(lora_dir))
    except Exception as e:
        print(f"[error] Failed to download dynamic LoRA from {repo}/{filename}: {e}", flush=True)
        raise RuntimeError(f"Failed to download dynamic LoRA: {e}")


def _add_loras_to_workflow(workflow, model_node_ref, lora_name, lora_strength, lora2_name=None, lora2_strength=0.0):
    current_model = model_node_ref

    if lora_name and lora_name != LORA_NONE and float(lora_strength) != 0.0:
        workflow["20"] = {
            "class_type": "LoraLoaderModelOnly",
            "inputs": {
                "model": current_model,
                "lora_name": lora_name,
                "strength_model": float(lora_strength),
            },
        }
        current_model = ["20", 0]

    if lora2_name and lora2_name != LORA_NONE and float(lora2_strength) != 0.0:
        workflow["21"] = {
            "class_type": "LoraLoaderModelOnly",
            "inputs": {
                "model": current_model,
                "lora_name": lora2_name,
                "strength_model": float(lora2_strength),
            },
        }
        current_model = ["21", 0]

    return current_model


def _build_workflow(
    prompt,
    negative_prompt,
    width,
    height,
    steps,
    cfg,
    seed,
    sampler,
    scheduler,
    lora_name=LORA_NONE,
    lora_strength=1.0,
    lora2_name=None,
    lora2_strength=0.0,
    unet_name=MODEL_FILE,
):
    workflow = {
        "1": {
            "class_type": "UNETLoader",
            "inputs": {"unet_name": unet_name, "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]},
        }

    model_ref = _add_loras_to_workflow(workflow, ["1", 0], lora_name, lora_strength, lora2_name, lora2_strength)
    workflow["5"]["inputs"]["model"] = model_ref
    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)
    
    try:
        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()
    finally:
        if temp_path.exists():
            temp_path.unlink()
            print(f"[ZDR] Deleted temp upload file from /tmp: {temp_path}", flush=True)

    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 _build_edit_workflow(
    input_filename,
    prompt,
    negative_prompt,
    steps,
    cfg,
    seed,
    sampler,
    scheduler,
    denoise,
    lora_name=LORA_NONE,
    lora_strength=1.0,
    lora2_name=None,
    lora2_strength=0.0,
    unet_name=MODEL_FILE,
):
    workflow = {
        "1": {
            "class_type": "UNETLoader",
            "inputs": {"unet_name": unet_name, "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]},
        }

    model_ref = _add_loras_to_workflow(workflow, ["1", 0], lora_name, lora_strength, lora2_name, lora2_strength)
    workflow["5"]["inputs"]["model"] = model_ref
    return workflow


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", []):
            filename = image["filename"]
            subfolder = image.get("subfolder", "")
            image_type = image.get("type", "output")
            
            if image_type == "output":
                file_path = COMFY_DIR / "output" / subfolder / filename
            else:
                file_path = COMFY_DIR / "temp" / subfolder / filename
                
            if not file_path.exists():
                raise RuntimeError(f"Generated file not found on disk: {file_path}")
                
            file_bytes = file_path.read_bytes()
            
            img = Image.open(io.BytesIO(file_bytes))
            prompt_data = img.info.get("prompt", "")
            workflow_data = img.info.get("workflow", "")
            
            prompt_bytes = prompt_data.encode('utf-8') if isinstance(prompt_data, str) else prompt_data
            workflow_bytes = workflow_data.encode('utf-8') if isinstance(workflow_data, str) else workflow_data
            
            exif_dict = {
                "0th": {
                    piexif.ImageIFD.Make: prompt_bytes,
                    piexif.ImageIFD.ImageDescription: workflow_bytes
                }
            }
            exif_bytes = piexif.dump(exif_dict)
            
            buffer = io.BytesIO()
            img.convert("RGB").save(buffer, "WEBP", exif=exif_bytes, quality=90)
            webp_bytes = buffer.getvalue()
            
            base64_str = base64.b64encode(webp_bytes).decode("utf-8")
            data_url = f"data:image/webp;base64,{base64_str}"
            
            try:
                file_path.unlink()
                print(f"[ZDR] Generation file successfully deleted from Space disk: {file_path}", flush=True)
            except Exception as e:
                print(f"[ZDR] Failed to delete generation file {file_path}: {e}", flush=True)
                
            return data_url
            
    raise RuntimeError("ComfyUI completed without returning an image.")


def _clamp_duration(value):
    try:
        value = int(value)
    except (TypeError, ValueError):
        value = DURATION_DEFAULT
    return max(DURATION_MIN, min(DURATION_MAX, value))


def _duration(
    prompt,
    negative_prompt,
    width,
    height,
    steps,
    cfg,
    seed,
    randomize_seed,
    sampler,
    scheduler,
    lora_name,
    lora_strength,
    lora2_repo,
    lora2_file,
    lora2_strength,
    custom_model_repo,
    custom_model_file,
    gpu_duration,
):
    return _clamp_duration(gpu_duration)


def _edit_duration(
    input_image,
    prompt,
    negative_prompt,
    width,
    height,
    steps,
    cfg,
    denoise,
    seed,
    randomize_seed,
    sampler,
    scheduler,
    lora_name,
    lora_strength,
    edit_lora2_repo,
    edit_lora2_file,
    edit_lora2_strength,
    edit_custom_model_repo,
    edit_custom_model_file,
    gpu_duration,
):
    return _clamp_duration(gpu_duration)


@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",
    lora_name=LORA_NONE,
    lora_strength=1.0,
    lora2_repo="",
    lora2_file="",
    lora2_strength=0.0,
    custom_model_repo="",
    custom_model_file="",
    gpu_duration=DURATION_DEFAULT,
):
    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()
        unet_name = _download_dynamic_model(custom_model_repo, custom_model_file)
        lora2_name = _download_dynamic_lora(lora2_repo, lora2_file)
        workflow = _build_workflow(
            prompt,
            negative_prompt,
            width,
            height,
            steps,
            cfg,
            seed,
            sampler,
            scheduler,
            lora_name,
            lora_strength,
            lora2_name,
            lora2_strength,
            unet_name=unet_name,
        )
        prompt_id = _queue_prompt(workflow)
        history_item = _wait_for_history(prompt_id)
        data_url = _load_output_image(history_item)
        
        html_img = f'<img src="{data_url}" style="max-width:100%; max-height:520px; object-fit:contain; margin:auto; display:block; border-radius:8px;">'
        return html_img, 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",
    lora_name=LORA_NONE,
    lora_strength=1.0,
    edit_lora2_repo="",
    edit_lora2_file="",
    edit_lora2_strength=0.0,
    edit_custom_model_repo="",
    edit_custom_model_file="",
    gpu_duration=DURATION_DEFAULT,
):
    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()
        unet_name = _download_dynamic_model(edit_custom_model_repo, edit_custom_model_file)
        lora2_name = _download_dynamic_lora(edit_lora2_repo, edit_lora2_file)
        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,
            lora_name,
            lora_strength,
            lora2_name,
            lora2_strength,
            unet_name=unet_name,
        )
        prompt_id = _queue_prompt(workflow)
        history_item = _wait_for_history(prompt_id)
        data_url = _load_output_image(history_item)
        
        try:
            input_file_path = COMFY_DIR / "input" / input_filename
            if input_file_path.exists():
                input_file_path.unlink()
                print(f"[ZDR] Input file successfully deleted from Space disk: {input_file_path}", flush=True)
        except Exception as e:
            print(f"[ZDR] Failed to delete input file: {e}", flush=True)
            
        html_img = f'<img src="{data_url}" style="max-width:100%; max-height:520px; object-fit:contain; margin:auto; display:block; border-radius:8px;">'
        return html_img, seed
    except Exception as exc:
        raise gr.Error(str(exc)) from exc


def refresh_loras():
    choices = _list_loras()
    return gr.update(choices=choices, value=LORA_NONE)


CSS = """
.gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
#result-image { min-height: 520px; }

/* Temporary hotfix for Chromium 144+ subpixel layout shifting in HF iframes */
.gradio-container { display: inline-table !important; width: 0px !important; height: 0px !important; overflow: hidden !important; opacity: 0 !important; pointer-events: none !important; }
"""

DURATION_INFO = (
    "How many seconds of ZeroGPU quota to request for this call. Krea2-Turbo "
    "usually finishes in ~10-20s once ComfyUI is warm - keep this low to save "
    "quota and get better queue priority. Bump it up only if a call times out "
    "(e.g. the very first generation after the Space restarts, while ComfyUI "
    "and the model are still loading)."
)

LORA_INFO = (
    "Drop .safetensors files into ComfyUI/models/loras in the repo, then hit "
    "Refresh. 'none' skips LoRA entirely."
)

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():
                        width = gr.Slider(512, 3072, value=1024, step=64, label="Width")
                        height = gr.Slider(512, 3072, 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("Custom Checkpoint (HuggingFace Dynamic)", open=False):
                        custom_model_repo = gr.Textbox(label="HF Repo ID", placeholder="e.g., beznogim666/test2", value="")
                        custom_model_file = gr.Textbox(label="Filename", placeholder="e.g., test_rc3.safetensors", value="")

                    with gr.Row():
                        lora_name = gr.Dropdown(
                            choices=[LORA_NONE, LORA_FILE],
                            value=LORA_FILE,
                            label="LoRA 1",
                            info=LORA_INFO,
                            allow_custom_value=True,
                            scale=4,
                        )
                        lora_refresh = gr.Button("Refresh", scale=1)
                    lora_strength = gr.Slider(0.0, 10.0, value=0.0, step=0.05, label="LoRA 1 strength")
                    
                    with gr.Accordion("LoRA 2 (HuggingFace Dynamic)", open=False):
                        lora2_repo = gr.Textbox(label="HF Repo ID", placeholder="e.g., beznogim666/test2", value="")
                        lora2_file = gr.Textbox(label="Filename", placeholder="e.g., lora2.safetensors", value="")
                        lora2_strength = gr.Slider(0.0, 10.0, value=0.0, step=0.05, label="LoRA 2 strength")

                    gpu_duration = gr.Slider(
                        DURATION_MIN,
                        DURATION_MAX,
                        value=DURATION_DEFAULT,
                        step=5,
                        label="ZeroGPU duration (s)",
                        info=DURATION_INFO,
                    )
                    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.HTML(label="Result", elem_id="result-image")

            inputs = [
                prompt,
                negative_prompt,
                width,
                height,
                steps,
                cfg,
                seed,
                randomize_seed,
                sampler,
                scheduler,
                lora_name,
                lora_strength,
                lora2_repo,
                lora2_file,
                lora2_strength,
                custom_model_repo,
                custom_model_file,
                gpu_duration,
            ]
            run.click(generate, inputs, [output, seed])
            prompt.submit(generate, inputs, [output, seed])
            lora_refresh.click(refresh_loras, None, lora_name)

        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, 3072, value=1024, step=64, label="Width")
                        edit_height = gr.Slider(512, 3072, 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("Custom Checkpoint (HuggingFace Dynamic)", open=False):
                        edit_custom_model_repo = gr.Textbox(label="HF Repo ID", placeholder="e.g., beznogim666/test2", value="")
                        edit_custom_model_file = gr.Textbox(label="Filename", placeholder="e.g., test_rc3.safetensors", value="")

                    with gr.Row():
                        edit_lora_name = gr.Dropdown(
                            choices=[LORA_NONE, LORA_FILE],
                            value=LORA_FILE,
                            label="LoRA 1",
                            info=LORA_INFO,
                            allow_custom_value=True,
                            scale=4,
                        )
                        edit_lora_refresh = gr.Button("Refresh", scale=1)
                    edit_lora_strength = gr.Slider(-2.0, 2.0, value=1.0, step=0.05, label="LoRA 1 strength")
                    
                    with gr.Accordion("LoRA 2 (HuggingFace Dynamic)", open=False):
                        edit_lora2_repo = gr.Textbox(label="HF Repo ID", placeholder="e.g., beznogim666/test2", value="")
                        edit_lora2_file = gr.Textbox(label="Filename", placeholder="e.g., lora2.safetensors", value="")
                        edit_lora2_strength = gr.Slider(-2.0, 2.0, value=0.0, step=0.05, label="LoRA 2 strength")

                    edit_gpu_duration = gr.Slider(
                        DURATION_MIN,
                        DURATION_MAX,
                        value=DURATION_DEFAULT,
                        step=5,
                        label="ZeroGPU duration (s)",
                        info=DURATION_INFO,
                    )
                    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.HTML(label="Edited image", 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_lora_name,
                edit_lora_strength,
                edit_lora2_repo,
                edit_lora2_file,
                edit_lora2_strength,
                edit_custom_model_repo,
                edit_custom_model_file,
                edit_gpu_duration,
            ]
            edit_run.click(edit_image, edit_inputs, [edit_output, edit_seed])
            edit_prompt.submit(edit_image, edit_inputs, [edit_output, edit_seed])
            edit_lora_refresh.click(refresh_loras, None, edit_lora_name)


_ensure_comfyui()

if __name__ == "__main__":
    demo.queue().launch()