Spaces:
Runtime error
Runtime error
| #!/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() | |