import os, json, requests, runpod import torch import random import comfy from comfy.sd import load_checkpoint_guess_config import nodes from nodes import NODE_CLASS_MAPPINGS from comfy_extras import nodes_post_processing, nodes_differential_diffusion, nodes_upscale_model import numpy as np from PIL import Image import asyncio import execution import server from nodes import load_custom_node from math import ceil, floor def download_file(url, save_dir='/content/ComfyUI/input'): os.makedirs(save_dir, exist_ok=True) file_name = url.split('/')[-1] file_path = os.path.join(save_dir, file_name) response = requests.get(url) response.raise_for_status() with open(file_path, 'wb') as file: file.write(response.content) return file_path loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) server_instance = server.PromptServer(loop) execution.PromptQueue(server_instance) load_custom_node("/content/ComfyUI/custom_nodes/ComfyUI-AutomaticCFG") load_custom_node("/content/ComfyUI/custom_nodes/ComfyUI-Custom-Scripts") load_custom_node("/content/ComfyUI/custom_nodes/Derfuu_ComfyUI_ModdedNodes") load_custom_node("/content/ComfyUI/custom_nodes/ComfyUI-Impact-Pack") load_custom_node("/content/ComfyUI/custom_nodes/ComfyUI-Inspire-Pack") load_custom_node("/content/ComfyUI/custom_nodes/ComfyUI-KJNodes") load_custom_node("/content/ComfyUI/custom_nodes/comfyui_controlnet_aux") load_custom_node("/content/ComfyUI/custom_nodes/ComfyUI-TiledDiffusion") load_custom_node("/content/ComfyUI/custom_nodes/was-node-suite-comfyui") Automatic_CFG = NODE_CLASS_MAPPINGS["Automatic CFG"]() ImageScaleToTotalPixels = nodes_post_processing.NODE_CLASS_MAPPINGS["ImageScaleToTotalPixels"]() GetImageSizeAndCount = NODE_CLASS_MAPPINGS["GetImageSizeAndCount"]() TTPlanet_TileSimple = NODE_CLASS_MAPPINGS["TTPlanet_TileSimple_Preprocessor"]() TiledDiffusion = NODE_CLASS_MAPPINGS["TiledDiffusion"]() KSampler_inspire = NODE_CLASS_MAPPINGS["KSampler //Inspire"]() ControlNetApplyAdvanced = NODE_CLASS_MAPPINGS["ControlNetApplyAdvanced"]() UltralyticsDetectorProvider = NODE_CLASS_MAPPINGS["UltralyticsDetectorProvider"]() SegmDetectorSEGS = NODE_CLASS_MAPPINGS["SegmDetectorSEGS"]() DifferentialDiffusion = nodes_differential_diffusion.NODE_CLASS_MAPPINGS["DifferentialDiffusion"]() DetailerForEach = NODE_CLASS_MAPPINGS["DetailerForEach"]() VAEDecodeTiled = NODE_CLASS_MAPPINGS["VAEDecodeTiled"]() ColorMatch = NODE_CLASS_MAPPINGS["ColorMatch"]() ImageBlend = nodes_post_processing.NODE_CLASS_MAPPINGS["ImageBlend"]() WAS_Image_Blending_Mode = NODE_CLASS_MAPPINGS["Image Blending Mode"]() ImageScale = NODE_CLASS_MAPPINGS["ImageScale"]() ImageScaleBy = NODE_CLASS_MAPPINGS["ImageScaleBy"]() UpscaleModelLoader = nodes_upscale_model.NODE_CLASS_MAPPINGS["UpscaleModelLoader"]() ImageUpscaleWithModel = nodes_upscale_model.NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]() with torch.inference_mode(): model_patcher, clip, vae, clipvision = load_checkpoint_guess_config("/content/ComfyUI/models/checkpoints/dreamshaperXL_lightningDPMSDE.safetensors", output_vae=True, output_clip=True, embedding_directory=None) tile_control_net = comfy.controlnet.load_controlnet("/content/ComfyUI/models/controlnet/xinsir-controlnet-tile-sdxl-1.0.safetensors") segm_detector = UltralyticsDetectorProvider.doit(model_name="segm/PitEyeDetailer-v2-seg.pt") upscale_model = UpscaleModelLoader.load_model(model_name="4xRealWebPhoto_v4_dat2.safetensors")[0] model_patcher = Automatic_CFG.patch(model=model_patcher, hard_mode=True, boost=True)[0] @torch.inference_mode() def generate(input): values = input["input"] input_image = values['input_image_check'] input_image = download_file(input_image) positive_prompt = values['positive_prompt'] negative_prompt = values['negative_prompt'] inspire_seed = values['inspire_seed'] inspire_steps = values['inspire_steps'] inspire_cfg = values['inspire_cfg'] inspire_sampler_name = values['inspire_sampler_name'] inspire_scheduler = values['inspire_scheduler'] inspire_denoise = values['inspire_denoise'] inspire_noise_mode = values['inspire_noise_mode'] inspire_batch_seed_mode = values['inspire_batch_seed_mode'] inspire_variation_seed = values['inspire_variation_seed'] inspire_variation_strength = values['inspire_variation_strength'] inspire_variation_method = values['inspire_variation_method'] scale_factor = values['scale_factor'] blur_strength = values['blur_strength'] strength = values['strength'] start_percent = values['start_percent'] end_percent = values['end_percent'] tile_method = values['tile_method'] tile_overlap = values['tile_overlap'] tile_size = values['tile_size'] threshold = values['threshold'] dilation = values['dilation'] crop_factor = values['crop_factor'] drop_size = values['drop_size'] labels = values['labels'] detailer_guide_size = values['detailer_guide_size'] detailer_guide_size_for_bbox = values['detailer_guide_size_for_bbox'] detailer_max_size = values['detailer_max_size'] detailer_seed = values['detailer_seed'] detailer_steps = values['detailer_steps'] detailer_cfg = values['detailer_cfg'] detailer_sampler_name = values['detailer_sampler_name'] detailer_scheduler = values['detailer_scheduler'] detailer_denoise = values['detailer_denoise'] detailer_feather = values['detailer_feather'] detailer_noise_mask = values['detailer_noise_mask'] detailer_force_inpaint = values['detailer_force_inpaint'] detailer_cycle = values['detailer_cycle'] detailer_inpaint_model = values['detailer_inpaint_model'] detailer_noise_mask_feather = values['detailer_noise_mask_feather'] color_method = values['color_method'] blend_factor = values['blend_factor'] blend_mode = values['blend_mode'] blending_mode = values['blending_mode'] blending_blend_percentage = values['blending_blend_percentage'] vram = values['vram'] upscale_mp = values['upscale_mp'] w_tiles = values['w_tiles'] h_tiles = values['h_tiles'] downscale_by = values['downscale_by'] output_image, output_mask = nodes.LoadImage().load_image(input_image) output_image_s = ImageScaleToTotalPixels.upscale(image=output_image, upscale_method="nearest-exact", megapixels=1.0)[0] image_width = GetImageSizeAndCount.getsize(output_image_s)["result"][1] image_height = GetImageSizeAndCount.getsize(output_image_s)["result"][2] w_math = ceil((image_width * upscale_mp) / 8) * 8 h_math = ceil((image_height * upscale_mp) / 8) * 8 tile_width = ceil((w_math / w_tiles) / 8) * 8 tile_height = ceil((h_math / h_tiles) / 8) * 8 tile_batch_size = floor((vram-3) / ((tile_width*tile_height) / 1000000)) upscale_image = ImageScaleBy.upscale(image=output_image, upscale_method="bilinear", scale_by=downscale_by)[0] upscaled_image = ImageUpscaleWithModel.upscale(upscale_model=upscale_model, image=upscale_image)[0] output_image = ImageScale.upscale(image=upscaled_image, upscale_method="bilinear", width=w_math, height=h_math, crop="disabled")[0] cond, pooled = clip.encode_from_tokens(clip.tokenize(positive_prompt), return_pooled=True) cond = [[cond, {"pooled_output": pooled}]] n_cond, n_pooled = clip.encode_from_tokens(clip.tokenize(negative_prompt), return_pooled=True) n_cond = [[n_cond, {"pooled_output": n_pooled}]] output_image_t = TTPlanet_TileSimple.execute(output_image, scale_factor=scale_factor, blur_strength=blur_strength)[0] positive, negative = ControlNetApplyAdvanced.apply_controlnet(positive=cond, negative=n_cond, control_net=tile_control_net, image=output_image_t, strength=strength, start_percent=start_percent, end_percent=end_percent) tile_model = TiledDiffusion.apply(model=model_patcher, method=tile_method, tile_width=tile_width, tile_height=tile_height, tile_overlap=tile_overlap, tile_batch_size=tile_batch_size)[0] latent_image = nodes.VAEEncode().encode(vae, output_image)[0] inspire_sample = KSampler_inspire.doit(model=tile_model, seed=inspire_seed, steps=inspire_steps, cfg=inspire_cfg, sampler_name=inspire_sampler_name, scheduler=inspire_scheduler, positive=positive, negative=negative, latent_image=latent_image, denoise=inspire_denoise, noise_mode=inspire_noise_mode, batch_seed_mode=inspire_batch_seed_mode, variation_seed=inspire_variation_seed, variation_strength=inspire_variation_strength, variation_method=inspire_variation_method)[0] tiled_decoded = VAEDecodeTiled.decode(vae=vae, samples=inspire_sample, tile_size=tile_size)[0] segs = SegmDetectorSEGS.doit(segm_detector=segm_detector[1], image=output_image, threshold=threshold, dilation=dilation, crop_factor=crop_factor, drop_size=drop_size, labels=labels)[0] dd_model_patcher = DifferentialDiffusion.apply(model_patcher)[0] detailer_image = DetailerForEach.do_detail(image=tiled_decoded, segs=segs, model=dd_model_patcher, clip=clip, vae=vae, guide_size=detailer_guide_size, guide_size_for_bbox=detailer_guide_size_for_bbox, max_size=detailer_max_size, seed=detailer_seed, steps=detailer_steps, cfg=detailer_cfg, sampler_name=detailer_sampler_name, scheduler=detailer_scheduler, positive=cond, negative=n_cond, denoise=detailer_denoise, feather=detailer_feather, noise_mask=detailer_noise_mask, force_inpaint=detailer_force_inpaint, cycle=detailer_cycle, inpaint_model=detailer_inpaint_model, noise_mask_feather=detailer_noise_mask_feather)[0] color_image = ColorMatch.colormatch(image_ref=output_image, image_target=detailer_image, method=color_method)[0] blend_image = ImageBlend.blend_images(image1=color_image, image2=detailer_image, blend_factor=blend_factor, blend_mode=blend_mode)[0] blending_image = WAS_Image_Blending_Mode.image_blending_mode(image_a=blend_image, image_b=output_image, mode=blending_mode, blend_percentage=blending_blend_percentage)[0] Image.fromarray(np.array(blending_image*255, dtype=np.uint8)[0]).save("/content/ultralytics.png") result = "/content/ultralytics.png" try: notify_uri = values['notify_uri'] del values['notify_uri'] notify_token = values['notify_token'] del values['notify_token'] discord_id = values['discord_id'] del values['discord_id'] if(discord_id == "discord_id"): discord_id = os.getenv('com_camenduru_discord_id') discord_channel = values['discord_channel'] del values['discord_channel'] if(discord_channel == "discord_channel"): discord_channel = os.getenv('com_camenduru_discord_channel') discord_token = values['discord_token'] del values['discord_token'] if(discord_token == "discord_token"): discord_token = os.getenv('com_camenduru_discord_token') job_id = values['job_id'] del values['job_id'] default_filename = os.path.basename(result) with open(result, "rb") as file: files = {default_filename: file.read()} payload = {"content": f"{json.dumps(values)} <@{discord_id}>"} response = requests.post( f"https://discord.com/api/v9/channels/{discord_channel}/messages", data=payload, headers={"Authorization": f"Bot {discord_token}"}, files=files ) response.raise_for_status() result_url = response.json()['attachments'][0]['url'] notify_payload = {"jobId": job_id, "result": result_url, "status": "DONE"} web_notify_uri = os.getenv('com_camenduru_web_notify_uri') web_notify_token = os.getenv('com_camenduru_web_notify_token') if(notify_uri == "notify_uri"): requests.post(web_notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token}) else: requests.post(web_notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token}) requests.post(notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": notify_token}) return {"jobId": job_id, "result": result_url, "status": "DONE"} except Exception as e: error_payload = {"jobId": job_id, "status": "FAILED"} try: if(notify_uri == "notify_uri"): requests.post(web_notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token}) else: requests.post(web_notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token}) requests.post(notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": notify_token}) except: pass return {"jobId": job_id, "result": f"FAILED: {str(e)}", "status": "FAILED"} finally: if os.path.exists(result): os.remove(result) runpod.serverless.start({"handler": generate})