#!/usr/bin/env python3 """ Klein Style Transfer - Gradio UI Klein LoRA + Dual Control (1 image + LoRA + DWPose + Depth) """ import os import sys import time import json import base64 import requests import tempfile import traceback import gradio as gr sys.stdout.reconfigure(line_buffering=True) # ============================================================================= # API CONFIG # ============================================================================= RUNPOD_API_KEY = os.environ.get('RUNPOD_API_KEY', '') ENDPOINT_ID = os.environ.get('RUNPOD_ENDPOINT_ID', '') BASE_URL = f"https://api.runpod.ai/v2/{ENDPOINT_ID}" HEADERS = { "Authorization": f"Bearer {RUNPOD_API_KEY}", "Content-Type": "application/json", } # Tinify upload config UPLOAD_ENDPOINT = os.environ.get('UPLOAD_ENDPOINT', '') TINIFY_API_KEY = os.environ.get('TINIFY_API_KEY', '') # ============================================================================= # WORKFLOW # ============================================================================= def build_workflow(prompt, seed, steps, lora_name, lora_strength_model, lora_strength_clip): """Klein LoRA + Dual ControlNet (DWPose + Depth).""" return { "6": { "inputs": {"image": "placeholder.png"}, "class_type": "LoadImage", "_meta": {"title": "Load Image (Target)"} }, "10": { "inputs": {"preprocessor": "DWPreprocessor", "resolution": 768, "image": ["6", 0]}, "class_type": "AIO_Preprocessor", "_meta": {"title": "AIO Aux Preprocessor (DWPose)"} }, "15": { "inputs": {"preprocessor": "DepthAnythingV2Preprocessor", "resolution": 768, "image": ["6", 0]}, "class_type": "AIO_Preprocessor", "_meta": {"title": "AIO Aux Preprocessor (Depth)"} }, "11": { "inputs": {"filename_prefix": "ComfyUI", "images": ["9:65", 0]}, "class_type": "SaveImage", "_meta": {"title": "Save Image"} }, "14": { "inputs": { "lora_name": lora_name, "strength_model": lora_strength_model, "strength_clip": lora_strength_clip, "model": ["9:70", 0], "clip": ["9:71", 0] }, "class_type": "LoraLoader", "_meta": {"title": "Load LoRA (Model and CLIP)"} }, "9:72": { "inputs": {"vae_name": "flux2-vae.safetensors"}, "class_type": "VAELoader", "_meta": {"title": "Load VAE"} }, "9:70": { "inputs": {"unet_name": "flux-2-klein-9b.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader", "_meta": {"title": "Load Diffusion Model"} }, "9:71": { "inputs": {"clip_name": "qwen_3_8b_fp8mixed.safetensors", "type": "flux2", "device": "default"}, "class_type": "CLIPLoader", "_meta": {"title": "Load CLIP"} }, "9:85": { "inputs": {"upscale_method": "nearest-exact", "megapixels": 1, "resolution_steps": 1, "image": ["6", 0]}, "class_type": "ImageScaleToTotalPixels", "_meta": {"title": "ImageScaleToTotalPixels (Target)"} }, "9:80": { "inputs": {"upscale_method": "nearest-exact", "megapixels": 1, "resolution_steps": 1, "image": ["10", 0]}, "class_type": "ImageScaleToTotalPixels", "_meta": {"title": "ImageScaleToTotalPixels (DWPose)"} }, "9:82": { "inputs": {"upscale_method": "nearest-exact", "megapixels": 1, "resolution_steps": 1, "image": ["15", 0]}, "class_type": "ImageScaleToTotalPixels", "_meta": {"title": "ImageScaleToTotalPixels (Depth)"} }, "9:74": { "inputs": {"text": prompt, "clip": ["14", 1]}, "class_type": "CLIPTextEncode", "_meta": {"title": "CLIP Text Encode (Positive Prompt)"} }, "9:87": { "inputs": {"text": "", "clip": ["14", 1]}, "class_type": "CLIPTextEncode", "_meta": {"title": "CLIP Text Encode ( Negative Prompt)"} }, "9:81": { "inputs": {"image": ["9:85", 0]}, "class_type": "GetImageSize", "_meta": {"title": "Get Image Size"} }, "9:79:78": { "inputs": {"pixels": ["9:80", 0], "vae": ["9:72", 0]}, "class_type": "VAEEncode", "_meta": {"title": "VAE Encode (DWPose)"} }, "9:83:78": { "inputs": {"pixels": ["9:82", 0], "vae": ["9:72", 0]}, "class_type": "VAEEncode", "_meta": {"title": "VAE Encode (Depth)"} }, "9:84:78": { "inputs": {"pixels": ["9:85", 0], "vae": ["9:72", 0]}, "class_type": "VAEEncode", "_meta": {"title": "VAE Encode (Target)"} }, "9:79:77": { "inputs": {"conditioning": ["9:74", 0], "latent": ["9:79:78", 0]}, "class_type": "ReferenceLatent", "_meta": {"title": "ReferenceLatent (DWPose + Positive)"} }, "9:79:76": { "inputs": {"conditioning": ["9:87", 0], "latent": ["9:79:78", 0]}, "class_type": "ReferenceLatent", "_meta": {"title": "ReferenceLatent (DWPose + Negative)"} }, "9:83:77": { "inputs": {"conditioning": ["9:79:77", 0], "latent": ["9:83:78", 0]}, "class_type": "ReferenceLatent", "_meta": {"title": "ReferenceLatent (Depth + Positive)"} }, "9:83:76": { "inputs": {"conditioning": ["9:79:76", 0], "latent": ["9:83:78", 0]}, "class_type": "ReferenceLatent", "_meta": {"title": "ReferenceLatent (Depth + Negative)"} }, "9:84:77": { "inputs": {"conditioning": ["9:83:77", 0], "latent": ["9:84:78", 0]}, "class_type": "ReferenceLatent", "_meta": {"title": "ReferenceLatent (Target + Positive)"} }, "9:84:76": { "inputs": {"conditioning": ["9:83:76", 0], "latent": ["9:84:78", 0]}, "class_type": "ReferenceLatent", "_meta": {"title": "ReferenceLatent (Target + Negative)"} }, "9:66": { "inputs": {"width": ["9:81", 0], "height": ["9:81", 1], "batch_size": 1}, "class_type": "EmptyFlux2LatentImage", "_meta": {"title": "Empty Flux 2 Latent"} }, "9:90": { "inputs": { "seed": seed, "steps": steps, "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1, "model": ["14", 0], "positive": ["9:84:77", 0], "negative": ["9:84:76", 0], "latent_image": ["9:66", 0] }, "class_type": "KSampler", "_meta": {"title": "KSampler"} }, "9:65": { "inputs": {"samples": ["9:90", 0], "vae": ["9:72", 0]}, "class_type": "VAEDecode", "_meta": {"title": "VAE Decode"} }, } # ============================================================================= # API HELPERS # ============================================================================= def upload_image(image_path): """Upload image to Tinify endpoint and return the public URL.""" file_name = os.path.basename(image_path) headers = {} if TINIFY_API_KEY: headers['x-api-key'] = TINIFY_API_KEY with open(image_path, 'rb') as f: files = {'file': (file_name, f, 'image/png')} response = requests.post(UPLOAD_ENDPOINT, files=files, headers=headers, timeout=120) response.raise_for_status() json_resp = response.json() for key in ['s3_url', 'url', 'link', 'image_url', 'file_url', 'imageUrl', 'fileUrl']: if key in json_resp: print(f"[Gradio] Uploaded: {json_resp[key]}") return json_resp[key] raise Exception(f"No URL in upload response: {json_resp}") def submit_and_wait(payload): """Submit job and poll until done. Returns (image_path, status_text).""" try: r = requests.post(f"{BASE_URL}/run", headers=HEADERS, json={"input": payload}, timeout=120) if r.status_code != 200: return None, f"Submit failed: HTTP {r.status_code} - {r.text[:500]}" resp = r.json() job_id = resp.get("id") if not job_id: return None, f"No job ID returned: {resp}" print(f"[Gradio] Job submitted: {job_id}") start = time.time() while (time.time() - start) < 600: time.sleep(3) elapsed = int(time.time() - start) try: status_resp = requests.get(f"{BASE_URL}/status/{job_id}", headers=HEADERS, timeout=60).json() except Exception as e: print(f"[Gradio] Poll error: {e}") continue status = status_resp.get("status", "UNKNOWN") print(f"[Gradio] [{elapsed}s] {status}") if status == "COMPLETED": output = status_resp.get("output", {}) if "error" in output: return None, f"ComfyUI error: {json.dumps(output['error'], indent=2)}" images = output.get("images", []) if images: img = images[0] if isinstance(img, dict): img = img.get('url', img.get('base64', '')) tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False) if isinstance(img, str) and img.startswith('http'): print(f"[Gradio] Downloading from: {img}") dl = requests.get(img, timeout=120) dl.raise_for_status() tmp.write(dl.content) else: if ',' in img: img = img.split(',')[1] tmp.write(base64.b64decode(img)) tmp.close() return tmp.name, f"Completed in {elapsed}s | Job: {job_id}" return None, f"No images in output: {json.dumps(output, indent=2)[:500]}" elif status in ("FAILED", "CANCELLED", "TIMED_OUT"): error = status_resp.get("error", status) output = status_resp.get("output", {}) error_detail = "" if isinstance(error, dict): error_detail = json.dumps(error, indent=2) elif isinstance(output, dict) and output: error_detail = json.dumps(output, indent=2) else: error_detail = str(error) return None, f"Failed ({status}):\n{error_detail[:1000]}" return None, "Timeout (600s)" except Exception as e: print(f"[Gradio] Exception: {traceback.format_exc()}") return None, f"Error: {str(e)}" # ============================================================================= # HANDLER # ============================================================================= def generate(target_image, prompt, steps, seed, lora_name, lora_strength_model, lora_strength_clip): """Handler for Klein LoRA + Dual ControlNet.""" if target_image is None: return None, "Please upload a target image" if not RUNPOD_API_KEY or not ENDPOINT_ID: return None, "Error: RUNPOD_API_KEY and RUNPOD_ENDPOINT_ID secrets are not set." if not UPLOAD_ENDPOINT: return None, "Error: UPLOAD_ENDPOINT secret is not set." try: print(f"[Gradio] Uploading image to Tinify...") image_url = upload_image(target_image) seed_val = int(seed) if int(seed) != -1 else int(time.time() * 1000) % (2**32) print(f"[Gradio] Building workflow...") workflow = build_workflow( prompt, seed_val, int(steps), lora_name, float(lora_strength_model), float(lora_strength_clip) ) payload = { "workflow": "custom", "payload": workflow, "params": {"node_6_image": image_url}, } print(f"[Gradio] Submitting job...") result_path, status = submit_and_wait(payload) return result_path, status except Exception as e: print(f"[Gradio] Handler error: {traceback.format_exc()}") return None, f"Error: {str(e)}" # ============================================================================= # GRADIO UI # ============================================================================= with gr.Blocks( title="Klein Style Transfer", theme=gr.themes.Soft(), css=""" .main-title { text-align: center; margin-bottom: 5px; } .sub-title { text-align: center; color: #666; margin-bottom: 20px; } .status-box textarea { font-family: monospace !important; font-size: 12px !important; } """ ) as demo: gr.Markdown("# Klein Style Transfer", elem_classes="main-title") gr.Markdown("LoRA + Dual ControlNet (DWPose + DepthAnythingV2) | Flux-2-Klein-9B", elem_classes="sub-title") with gr.Row(): with gr.Column(scale=1): target_image = gr.Image(label="Upload Target Image", type="filepath", height=300) prompt = gr.Textbox( label="Prompt", value="change this image into expressionist_style", lines=2, placeholder="Describe the style you want..." ) lora_name = gr.Dropdown( label="LoRA Style", choices=[ "expression_stylist_flux_lora.safetensors", "flux_klein_impasto_lora.safetensors", ], value="expression_stylist_flux_lora.safetensors", allow_custom_value=True ) with gr.Row(): lora_str_model = gr.Slider( minimum=0, maximum=2, value=0.8, step=0.05, label="LoRA Strength (Model)" ) lora_str_clip = gr.Slider( minimum=0, maximum=2, value=1.0, step=0.05, label="LoRA Strength (CLIP)" ) with gr.Row(): steps = gr.Slider(minimum=1, maximum=20, value=4, step=1, label="Steps") seed = gr.Number(value=42, label="Seed (-1 = random)", precision=0) generate_btn = gr.Button("Generate", variant="primary") with gr.Column(scale=1): output_image = gr.Image(label="Output", type="filepath", height=300) status_text = gr.Textbox( label="Status", interactive=False, lines=3, elem_classes="status-box" ) generate_btn.click( fn=generate, inputs=[target_image, prompt, steps, seed, lora_name, lora_str_model, lora_str_clip], outputs=[output_image, status_text], ) gr.Markdown("---") gr.Markdown( "Controls: DWPose + DepthAnythingV2 | " "Model: Flux-2-Klein-9B | Powered by RunPod Serverless" ) if __name__ == "__main__": demo.queue().launch()