computerwiz8432/IllusionDiffusion
0
1import torch2import gradio as gr3from gradio import processing_utils, utils4from PIL import Image5import random6from diffusers import (7 DiffusionPipeline,8 AutoencoderKL,9 StableDiffusionControlNetPipeline,10 ControlNetModel,11 StableDiffusionLatentUpscalePipeline,12 StableDiffusionImg2ImgPipeline,13 StableDiffusionControlNetImg2ImgPipeline,14 DPMSolverMultistepScheduler, # <-- Added import15 EulerDiscreteScheduler # <-- Added import16)17import time18from share_btn import community_icon_html, loading_icon_html, share_js19import user_history20from illusion_style import css21 22BASE_MODEL = "SG161222/Realistic_Vision_V5.1_noVAE"23 24# Initialize both pipelines25vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)26#init_pipe = DiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V5.1_noVAE", torch_dtype=torch.float16)27controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=torch.float16)#, torch_dtype=torch.float16)28main_pipe = StableDiffusionControlNetPipeline.from_pretrained(29 BASE_MODEL,30 controlnet=controlnet,31 vae=vae,32 safety_checker=None,33 torch_dtype=torch.float16,34).to("cuda")35 36#main_pipe.unet = torch.compile(main_pipe.unet, mode="reduce-overhead", fullgraph=True)37#main_pipe.unet.to(memory_format=torch.channels_last)38#main_pipe.unet = torch.compile(main_pipe.unet, mode="reduce-overhead", fullgraph=True)39#model_id = "stabilityai/sd-x2-latent-upscaler"40image_pipe = StableDiffusionControlNetImg2ImgPipeline(**main_pipe.components)41 42#image_pipe.unet = torch.compile(image_pipe.unet, mode="reduce-overhead", fullgraph=True)43#upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)44#upscaler.to("cuda")45 46 47# Sampler map48SAMPLER_MAP = {49 "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),50 "Euler": lambda config: EulerDiscreteScheduler.from_config(config),51}52 53def center_crop_resize(img, output_size=(512, 512)):54 width, height = img.size55 56 # Calculate dimensions to crop to the center57 new_dimension = min(width, height)58 left = (width - new_dimension)/259 top = (height - new_dimension)/260 right = (width + new_dimension)/261 bottom = (height + new_dimension)/262 63 # Crop and resize64 img = img.crop((left, top, right, bottom))65 img = img.resize(output_size)66 67 return img68 69def common_upscale(samples, width, height, upscale_method, crop=False):70 if crop == "center":71 old_width = samples.shape[3]72 old_height = samples.shape[2]73 old_aspect = old_width / old_height74 new_aspect = width / height75 x = 076 y = 077 if old_aspect > new_aspect:78 x = round((old_width - old_width * (new_aspect / old_aspect)) / 2)79 elif old_aspect < new_aspect:80 y = round((old_height - old_height * (old_aspect / new_aspect)) / 2)81 s = samples[:,:,y:old_height-y,x:old_width-x]82 else:83 s = samples84 85 return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)86 87def upscale(samples, upscale_method, scale_by):88 #s = samples.copy()89 width = round(samples["images"].shape[3] * scale_by)90 height = round(samples["images"].shape[2] * scale_by)91 s = common_upscale(samples["images"], width, height, upscale_method, "disabled")92 return (s)93 94def check_inputs(prompt: str, control_image: Image.Image):95 if control_image is None:96 raise gr.Error("Please select or upload an Input Illusion")97 if prompt is None or prompt == "":98 raise gr.Error("Prompt is required")99 100def convert_to_pil(base64_image):101 pil_image = processing_utils.decode_base64_to_image(base64_image)102 return pil_image103 104def convert_to_base64(pil_image):105 base64_image = processing_utils.encode_pil_to_base64(pil_image)106 return base64_image107 108# Inference function109def inference(110 control_image: Image.Image,111 prompt: str,112 negative_prompt: str,113 guidance_scale: float = 8.0,114 controlnet_conditioning_scale: float = 1,115 control_guidance_start: float = 1, 116 control_guidance_end: float = 1,117 upscaler_strength: float = 0.5,118 seed: int = -1,119 sampler = "DPM++ Karras SDE",120 progress = gr.Progress(track_tqdm=True),121 profile: gr.OAuthProfile | None = None,122):123 start_time = time.time()124 start_time_struct = time.localtime(start_time)125 start_time_formatted = time.strftime("%H:%M:%S", start_time_struct)126 print(f"Inference started at {start_time_formatted}")127 128 # Generate the initial image129 #init_image = init_pipe(prompt).images[0]130 131 # Rest of your existing code132 control_image_small = center_crop_resize(control_image)133 control_image_large = center_crop_resize(control_image, (1024, 1024))134 135 main_pipe.scheduler = SAMPLER_MAP[sampler](main_pipe.scheduler.config)136 my_seed = random.randint(0, 2**32 - 1) if seed == -1 else seed137 generator = torch.Generator(device="cuda").manual_seed(my_seed)138 139 out = main_pipe(140 prompt=prompt,141 negative_prompt=negative_prompt,142 image=control_image_small,143 guidance_scale=float(guidance_scale),144 controlnet_conditioning_scale=float(controlnet_conditioning_scale),145 generator=generator,146 control_guidance_start=float(control_guidance_start),147 control_guidance_end=float(control_guidance_end),148 num_inference_steps=15,149 output_type="latent"150 )151 upscaled_latents = upscale(out, "nearest-exact", 2)152 out_image = image_pipe(153 prompt=prompt,154 negative_prompt=negative_prompt,155 control_image=control_image_large, 156 image=upscaled_latents,157 guidance_scale=float(guidance_scale),158 generator=generator,159 num_inference_steps=20,160 strength=upscaler_strength,161 control_guidance_start=float(control_guidance_start),162 control_guidance_end=float(control_guidance_end),163 controlnet_conditioning_scale=float(controlnet_conditioning_scale)164 )165 end_time = time.time()166 end_time_struct = time.localtime(end_time)167 end_time_formatted = time.strftime("%H:%M:%S", end_time_struct)168 print(f"Inference ended at {end_time_formatted}, taking {end_time-start_time}s")169 170 # Save image + metadata171 user_history.save_image(172 label=prompt,173 image=out_image["images"][0],174 profile=profile,175 metadata={176 "prompt": prompt,177 "negative_prompt": negative_prompt,178 "guidance_scale": guidance_scale,179 "controlnet_conditioning_scale": controlnet_conditioning_scale,180 "control_guidance_start": control_guidance_start,181 "control_guidance_end": control_guidance_end,182 "upscaler_strength": upscaler_strength,183 "seed": seed,184 "sampler": sampler,185 },186 )187 188 return out_image["images"][0], gr.update(visible=True), gr.update(visible=True), my_seed189 190with gr.Blocks() as app:191 gr.Markdown(192 '''193 <center><h1>Illusion Diffusion HQ 🌀</h1></span> 194 <span font-size:16px;">Generate stunning high quality illusion artwork with Stable Diffusion</span> 195 </center>196 197 A space by AP [Follow me on Twitter](https://twitter.com/angrypenguinPNG) with big contributions from [multimodalart](https://twitter.com/multimodalart)198 199 This project works by using [Monster Labs QR Control Net](https://huggingface.co/monster-labs/control_v1p_sd15_qrcode_monster).200 Given a prompt and your pattern, we use a QR code conditioned controlnet to create a stunning illusion! Credit to: [MrUgleh](https://twitter.com/MrUgleh) for discovering the workflow :)201 '''202 )203 state_img_input = gr.State()204 state_img_output = gr.State()205 with gr.Row():206 with gr.Column():207 control_image = gr.Image(label="Input Illusion", type="pil", elem_id="control_image")208 controlnet_conditioning_scale = gr.Slider(minimum=0.0, maximum=5.0, step=0.01, value=0.8, label="Illusion strength", elem_id="illusion_strength", info="ControlNet conditioning scale")209 gr.Examples(examples=["checkers.png", "checkers_mid.jpg", "pattern.png", "ultra_checkers.png", "spiral.jpeg", "funky.jpeg" ], inputs=control_image)210 prompt = gr.Textbox(label="Prompt", elem_id="prompt", info="Type what you want to generate", placeholder="Medieval village scene with busy streets and castle in the distance")211 negative_prompt = gr.Textbox(label="Negative Prompt", info="Type what you don't want to see", value="low quality", elem_id="negative_prompt")212 with gr.Accordion(label="Advanced Options", open=False):213 guidance_scale = gr.Slider(minimum=0.0, maximum=50.0, step=0.25, value=7.5, label="Guidance Scale")214 sampler = gr.Dropdown(choices=list(SAMPLER_MAP.keys()), value="Euler")215 control_start = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0, label="Start of ControlNet")216 control_end = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1, label="End of ControlNet")217 strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1, label="Strength of the upscaler")218 seed = gr.Slider(minimum=-1, maximum=9999999999, step=1, value=-1, label="Seed", info="-1 means random seed")219 used_seed = gr.Number(label="Last seed used",interactive=False)220 run_btn = gr.Button("Run")221 with gr.Column():222 result_image = gr.Image(label="Illusion Diffusion Output", interactive=False, elem_id="output")223 with gr.Group(elem_id="share-btn-container", visible=False) as share_group:224 community_icon = gr.HTML(community_icon_html)225 loading_icon = gr.HTML(loading_icon_html)226 share_button = gr.Button("Share to community", elem_id="share-btn")227 228 prompt.submit(229 check_inputs,230 inputs=[prompt, control_image],231 queue=False232 ).success(233 convert_to_pil,234 inputs=[control_image],235 outputs=[state_img_input],236 queue=False,237 preprocess=False,238 ).success(239 inference,240 inputs=[state_img_input, prompt, negative_prompt, guidance_scale, controlnet_conditioning_scale, control_start, control_end, strength, seed, sampler],241 outputs=[state_img_output, result_image, share_group, used_seed]242 ).success(243 convert_to_base64,244 inputs=[state_img_output],245 outputs=[result_image],246 queue=False,247 postprocess=False248 )249 run_btn.click(250 check_inputs,251 inputs=[prompt, control_image],252 queue=False253 ).success(254 convert_to_pil,255 inputs=[control_image],256 outputs=[state_img_input],257 queue=False,258 preprocess=False,259 ).success(260 inference,261 inputs=[state_img_input, prompt, negative_prompt, guidance_scale, controlnet_conditioning_scale, control_start, control_end, strength, seed, sampler],262 outputs=[state_img_output, result_image, share_group, used_seed]263 ).success(264 convert_to_base64,265 inputs=[state_img_output],266 outputs=[result_image],267 queue=False,268 postprocess=False269 )270 share_button.click(None, [], [], _js=share_js)271 272with gr.Blocks(css=css) as app_with_history:273 with gr.Tab("Demo"):274 app.render()275 with gr.Tab("Past generations"):276 user_history.render()277 278app_with_history.queue(max_size=20,api_open=False )279 280if __name__ == "__main__":281 app_with_history.launch(max_threads=400)282