import os import json import time import uuid import shutil import subprocess import threading import requests from io import BytesIO from pathlib import Path # Load .env for local development (ignored if not installed / not present) try: from dotenv import load_dotenv load_dotenv() except ImportError: pass # Compatibility shim — newer huggingface_hub removed HfFolder which old gradio needs import huggingface_hub as _hfhub if not hasattr(_hfhub, "HfFolder"): class _HfFolder: @staticmethod def get_token(): return _hfhub.get_token() @staticmethod def save_token(token): pass _hfhub.HfFolder = _HfFolder import spaces import gradio as gr from PIL import Image from huggingface_hub import hf_hub_download # --------------------------------------------------------------------------- # Paths — use /data (persistent storage) if available, else /home/user # --------------------------------------------------------------------------- BASE_DIR = "/data" if os.path.exists("/data") else "/home/user" COMFYUI_DIR = f"{BASE_DIR}/ComfyUI" COMFYUI_INPUT = f"{COMFYUI_DIR}/input" COMFYUI_OUTPUT = f"{COMFYUI_DIR}/output" COMFYUI_MODELS = f"{COMFYUI_DIR}/models" COMFYUI_CUSTOM_NODES = f"{COMFYUI_DIR}/custom_nodes" COMFYUI_URL = "http://127.0.0.1:8188" # --------------------------------------------------------------------------- # Custom nodes # --------------------------------------------------------------------------- CUSTOM_NODES = { "ComfyUI-GGUF": "https://github.com/city96/ComfyUI-GGUF", "masquerade-nodes-comfyui": "https://github.com/BadCafeCode/masquerade-nodes-comfyui", "ComfyUI-KJNodes": "https://github.com/kijai/ComfyUI-KJNodes", "ComfyUI_LayerStyle_Advance": "https://github.com/chflame163/ComfyUI_LayerStyle_Advance", "comfyui-sam3": "https://github.com/PozzettiAndrea/ComfyUI-SAM3", } # --------------------------------------------------------------------------- # Models # --------------------------------------------------------------------------- MODELS = [ { "repo_id": "unsloth/FLUX.2-klein-4B-GGUF", "filename": "flux-2-klein-4b-Q8_0.gguf", "dest": f"{COMFYUI_MODELS}/unet/flux-2-klein-4b-Q8_0.gguf", }, { "repo_id": "Comfy-Org/z_image_turbo", "filename": "split_files/text_encoders/qwen_3_4b.safetensors", "dest": f"{COMFYUI_MODELS}/text_encoders/qwen_3_4b.safetensors", }, { "repo_id": "Comfy-Org/flux2-dev", "filename": "split_files/vae/flux2-vae.safetensors", "dest": f"{COMFYUI_MODELS}/vae/flux2-vae.safetensors", }, { "repo_id": "p1atdev/auraflow-v0.3-pvc-style-lora", "filename": "aura-pvc-2-_00010e_074520s.safetensors", "revision": "cafeee8ab8681ab679944b4e75ab0bdc4bdec6f7", "dest": f"{COMFYUI_MODELS}/loras/aura-pvc-2-_00010e_074520s.safetensors", }, ] # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def run(cmd: str, **kwargs): print(f"$ {cmd}") subprocess.run(cmd, shell=True, check=True, **kwargs) def download_model(model: dict): dest = Path(model["dest"]) if dest.exists(): print(f" already exists: {dest.name}") return dest.parent.mkdir(parents=True, exist_ok=True) print(f" downloading {dest.name} ...") kwargs = dict(repo_id=model["repo_id"], filename=model["filename"]) if "revision" in model: kwargs["revision"] = model["revision"] cached = hf_hub_download(**kwargs) shutil.copy(cached, dest) print(f" saved → {dest}") # --------------------------------------------------------------------------- # Setup — split into two parts: # setup_env() : clone repos, install packages, download models (no GPU needed) # start_comfyui(): launch ComfyUI subprocess (must run inside @spaces.GPU) # --------------------------------------------------------------------------- def setup_env(): # 1. Clone ComfyUI if not Path(f"{COMFYUI_DIR}/main.py").exists(): print("=== Cloning ComfyUI ===") run(f"git clone --depth 1 https://github.com/comfyanonymous/ComfyUI {COMFYUI_DIR}") print("=== Installing ComfyUI requirements ===") run(f"pip install -r {COMFYUI_DIR}/requirements.txt -q --break-system-packages") # 2. Custom nodes print("=== Installing custom nodes ===") os.makedirs(COMFYUI_CUSTOM_NODES, exist_ok=True) for name, url in CUSTOM_NODES.items(): node_dir = Path(f"{COMFYUI_CUSTOM_NODES}/{name}") if not node_dir.exists(): print(f" cloning {name}") run(f"git clone --depth 1 {url} {node_dir}") req = node_dir / "requirements.txt" if req.exists(): run(f"pip install -r {req} -q --break-system-packages") # 3. Models print("=== Downloading models ===") for model in MODELS: download_model(model) # 4. Config config_path = Path(f"{COMFYUI_DIR}/user/__manager/config.ini") if config_path.exists(): content = config_path.read_text() content = content.replace("network_mode = public", "network_mode = personal_cloud") content = content.replace("security_level = strict", "security_level = normal") config_path.write_text(content) _comfyui_started = False def start_comfyui(): """Start ComfyUI subprocess. Must be called from within a @spaces.GPU context.""" global _comfyui_started if _comfyui_started: return print("=== Starting ComfyUI ===") os.makedirs(COMFYUI_INPUT, exist_ok=True) os.makedirs(COMFYUI_OUTPUT, exist_ok=True) proc = subprocess.Popen( f"python {COMFYUI_DIR}/main.py --listen --port 8188", shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, ) def _stream(p): for line in p.stdout: print(line.decode(errors="replace"), end="") threading.Thread(target=_stream, args=(proc,), daemon=True).start() print("Waiting for ComfyUI...") for _ in range(60): try: if requests.get(f"{COMFYUI_URL}/system_stats", timeout=3).status_code == 200: print("ComfyUI ready!") _comfyui_started = True return except Exception: pass time.sleep(3) raise RuntimeError("ComfyUI failed to start within 3 minutes") # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- @spaces.GPU(duration=60) def generate(target_img, clothing_img, progress=gr.Progress(track_tqdm=True)): start_comfyui() if target_img is None or clothing_img is None: raise gr.Error("Please upload both images before generating.") # Save user images to ComfyUI input folder uid = uuid.uuid4().hex[:8] target_name = f"target_{uid}.png" clothing_name = f"clothing_{uid}.png" Image.fromarray(target_img).save(f"{COMFYUI_INPUT}/{target_name}") Image.fromarray(clothing_img).save(f"{COMFYUI_INPUT}/{clothing_name}") # Load workflow and inject filenames + fresh seed with open("workflow_api.json") as f: workflow = json.load(f) workflow["76"]["inputs"]["image"] = target_name workflow["81"]["inputs"]["image"] = clothing_name workflow["104"]["inputs"]["noise_seed"] = int(time.time() * 1000) % (2 ** 32) # Submit to ComfyUI progress(0.05, desc="Submitting to ComfyUI...") client_id = uuid.uuid4().hex resp = requests.post( f"{COMFYUI_URL}/prompt", json={"prompt": workflow, "client_id": client_id}, ) resp.raise_for_status() prompt_id = resp.json()["prompt_id"] # Poll /history until done progress(0.1, desc="Generating — this takes 1–2 minutes...") started = time.time() while True: history = requests.get(f"{COMFYUI_URL}/history/{prompt_id}").json() if prompt_id in history: entry = history[prompt_id] status = entry.get("status", {}) if status.get("status_str") == "error" or entry.get("error"): raise gr.Error(f"Generation failed: {entry.get('error', 'unknown error')}") break elapsed = int(time.time() - started) progress(min(0.9, 0.1 + elapsed / 150 * 0.8), desc=f"Generating... ({elapsed}s)") time.sleep(3) # Retrieve output image outputs = history[prompt_id].get("outputs", {}) for node_output in outputs.values(): if "images" in node_output: img_info = node_output["images"][0] img_bytes = requests.get( f"{COMFYUI_URL}/view", params={ "filename": img_info["filename"], "subfolder": img_info.get("subfolder", ""), "type": img_info.get("type", "output"), }, ).content progress(1.0, desc="Done!") return Image.open(BytesIO(img_bytes)) raise gr.Error("No output image was returned by ComfyUI.") # --------------------------------------------------------------------------- # Access code — set the same code somewhere visible in your CV # --------------------------------------------------------------------------- ACCESS_CODE = os.getenv("SECRET_CODE") CV_FILE = "Sofia_Metelitsa_CV.pdf" # --------------------------------------------------------------------------- # Run setup once at startup (downloads + custom nodes, no GPU needed) # ComfyUI itself is started on first generate() call when GPU is available # --------------------------------------------------------------------------- # Run the non-GPU parts now (clone repos, install packages, download models) setup_env() # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- with gr.Blocks(title="Cloth Swap", theme=gr.themes.Soft()) as demo: # ── Gate screen ────────────────────────────────────────────────────────── with gr.Column(visible=True) as gate: gr.Markdown( """ # Cloth Swap — Virtual Try-On *AI-powered outfit swapping built with FLUX.2 Kontext + ComfyUI* --- ### 🔑 To access this tool, download my CV. The access code will appear here after you download it. """ ) dl_btn = gr.DownloadButton( label="📄 Download CV", value=CV_FILE, variant="secondary", size="lg", ) code_reveal = gr.Markdown(visible=False) gr.Markdown("---") code_input = gr.Textbox( label="Enter access code from CV", placeholder="Access code", type="password", max_lines=1, ) unlock_btn = gr.Button("Unlock →", variant="primary") error_msg = gr.Markdown(visible=False) # ── Main app (hidden until unlocked) ───────────────────────────────────── with gr.Column(visible=False) as main_app: gr.Markdown( """ # Cloth Swap Upload a photo of a **person** and a **clothing reference**. The AI swaps the outfit while preserving pose, lighting, and expression. """ ) with gr.Row(): target_input = gr.Image(label="Person (target)", type="numpy", height=400) clothing_input = gr.Image(label="Clothing reference", type="numpy", height=400) btn = gr.Button("✨ Generate", variant="primary", size="lg") output = gr.Image(label="Result", height=500) gr.Markdown("*Generation takes ~1–2 minutes.*") btn.click( fn=generate, inputs=[target_input, clothing_input], outputs=output, ) # ── Logic ───────────────────────────────────────────────────────────────── def reveal_code(): """Called when user clicks Download CV — reveals the access code on screen.""" return gr.update(visible=True, value=f"**Your access code:** `{ACCESS_CODE}`") def check_code(entered): if entered.strip() == ACCESS_CODE: return ( gr.update(visible=False), # hide gate gr.update(visible=True), # show app gr.update(visible=False), # hide error ) return ( gr.update(visible=True), gr.update(visible=False), gr.update(visible=True, value="Incorrect code — check your CV and try again."), ) dl_btn.click(fn=reveal_code, outputs=[code_reveal]) unlock_btn.click( fn=check_code, inputs=[code_input], outputs=[gate, main_app, error_msg], ) # The HF base image ships websockets v13+ which removed the legacy API that # the installed uvicorn version uses. Redirect uvicorn's WebSocket backend to # wsproto before demo.launch() triggers the import of uvicorn.protocols.websockets.auto. import sys as _sys from types import ModuleType as _ModuleType _ws_auto = _ModuleType("uvicorn.protocols.websockets.auto") _sys.modules["uvicorn.protocols.websockets.auto"] = _ws_auto try: from uvicorn.protocols.websockets.wsproto_impl import WSProtocol as _WSP _ws_auto.AutoWebSocketsProtocol = _WSP print("uvicorn → wsproto WebSocket backend active") except Exception as _e: print(f"WARNING: wsproto backend setup failed: {_e}") # When server_name="0.0.0.0", gradio constructs local_url as "http://0.0.0.0:7860/" # which is a bind address and can't be used as a connection target — url_ok fails. # The server IS running; patch url_ok so gradio doesn't block on this false negative. try: import gradio.networking as _gnet _gnet.url_ok = lambda url: True except Exception as _e: print(f"url_ok patch failed: {_e}") # Starlette 0.36+ changed TemplateResponse(name, context) → TemplateResponse(request, name, context). # Gradio 4.44.0 still uses the old signature, so "index.html" ends up as `request` # and the context dict ends up as `name`. Patch to restore the old behaviour. try: import starlette.templating as _st _orig_TR = _st.Jinja2Templates.TemplateResponse def _compat_TR(self, *args, **kwargs): if args and isinstance(args[0], str): # Old-style call: TemplateResponse("template.html", context_dict, ...) name = args[0] context = args[1] if len(args) > 1 else kwargs.pop("context", {}) request = context.get("request") return _orig_TR(self, request, name, context, **kwargs) return _orig_TR(self, *args, **kwargs) _st.Jinja2Templates.TemplateResponse = _compat_TR except Exception as _e: print(f"starlette TemplateResponse patch failed: {_e}") # Jinja2 3.1.4 bug: LRUCache uses unhashable dict as cache key when globals is # non-empty, causing TemplateResponse to crash. Convert TypeError → KeyError so # templates are loaded fresh when the key can't be cached. try: import jinja2.utils as _jutils _LRU = _jutils.LRUCache _orig_gi = _LRU.__getitem__ _orig_si = _LRU.__setitem__ def _safe_gi(self, key): try: return _orig_gi(self, key) except TypeError: raise KeyError(key) def _safe_si(self, key, value): try: _orig_si(self, key, value) except TypeError: pass _LRU.__getitem__ = _safe_gi _LRU.__setitem__ = _safe_si except Exception as _e: print(f"jinja2 LRUCache patch failed: {_e}") # gradio_client bug: _json_schema_to_python_type() can't handle bool schemas # (e.g. additionalProperties: true). Patch the internal recursive function so # any non-dict schema is treated as "Any" instead of raising APIInfoParseError. try: import gradio_client.utils as _gcu _orig_j2p = _gcu._json_schema_to_python_type def _safe_j2p(schema, defs=None): if not isinstance(schema, dict): return "Any" return _orig_j2p(schema, defs) _gcu._json_schema_to_python_type = _safe_j2p except Exception as _e: print(f"gradio_client patch failed: {_e}") demo.launch(server_name="0.0.0.0", show_error=True, show_api=False)