import torch import gradio as gr from multiprocessing import cpu_count import os from generate_code import create_code, backgrounds, correction_map from diffusers import ( StableDiffusionControlNetPipeline, ControlNetModel, EulerAncestralDiscreteScheduler, ) main_generator = torch.Generator() # MONSTER_V2 = "/home/ubuntu/training/diffusers/examples/controlnet/out_model_2023-06-18_17-27-06" # LANDMARKS = "/home/ubuntu/training/diffusers/examples/controlnet/out_model_2023-06-19_23-43-50/" MONSTER_V2 = "monster-labs/V2" LANDMARKS = "monster-labs/V2" controlnet = [ ControlNetModel.from_pretrained(MONSTER_V2, torch_dtype=torch.float16, subfolder="step1", use_auth_token=os.environ["HUGGINGFACE_TOKEN"]), ControlNetModel.from_pretrained(LANDMARKS, torch_dtype=torch.float16, subfolder="step2", use_auth_token=os.environ["HUGGINGFACE_TOKEN"]), ] pipe = StableDiffusionControlNetPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16, ).to("cuda") pipe.enable_xformers_memory_efficient_attention() pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) def inference_map( qr_code_content: str, prompt: str, negative_prompt: str, guidance_scale: float, controlnet_conditioning_scale_0: float, seed: int, controlnet_start_0: float, controlnet_start_1: float, controlnet_end_0: float, controlnet_end_1: float, background: str, error_correction: str, margin: int, module_size: int, width: int, height: int, ): return inference( qr_code_content, prompt, negative_prompt, guidance_scale, (controlnet_conditioning_scale_0, .6 + (controlnet_conditioning_scale_0 * .8 / 2)), seed, (controlnet_start_0, controlnet_start_1), (controlnet_end_0, controlnet_end_1), background, error_correction, margin, module_size, width, height, ) def inference( qr_code_content: str, prompt: str, negative_prompt: str, guidance_scale: float = 10.0, controlnet_conditioning_scale: tuple[float, float] = (1.0, 1.0), seed: int = -1, controlnet_start: tuple[float, float] = (0.2, 0.0), controlnet_end: tuple[float, float] = (0.95, 1.0), background: str = "gray", error_correction: str = "H", margin: int = 1, module_size: int = 16, width: int = None, height: int = None, ): if prompt is None or prompt == "": raise gr.Error("Prompt is required") if qr_code_content is None or qr_code_content == "": raise gr.Error("QR Code Content is required") if background not in backgrounds: raise gr.Error("Invalid background") if error_correction not in correction_map: raise gr.Error("Invalid error correction") generator = torch.manual_seed(seed) if seed != -1 else main_generator # print("Generating QR Code from content") qrcode_images = create_code(qr_code_content, module_size, margin, background, error_correction, False, 1, True) out = pipe( prompt=prompt, negative_prompt=negative_prompt, image=list(qrcode_images), width=qrcode_images[0].width, height=qrcode_images[0].height, guidance_scale=float(guidance_scale), controlnet_conditioning_scale=controlnet_conditioning_scale, # controlnet_start=controlnet_start, # controlnet_end=controlnet_end, controlnet_guidance=[(controlnet_start[0], controlnet_end[0]), (controlnet_start[1], controlnet_end[1])], generator=generator, num_inference_steps=40, ) return out.images[0] with gr.Blocks() as blocks: with gr.Row(): with gr.Column(): qr_code_content = gr.Textbox( label="QR Code Content", info="QR Code Content or URL", value="", ) prompt = gr.Textbox( label="Prompt", info="Prompt that guides the generation towards", ) negative_prompt = gr.Textbox( label="Negative Prompt", value="ugly, disfigured, low quality, blurry, nsfw", ) with gr.Accordion( label="Params: The generated QR Code functionality is largely influenced by the parameters detailed below", open=True, ): controlnet_conditioning_scale = gr.Slider( minimum=0.5, maximum=2.5, step=0.01, value=1.5, label="Controlnet Conditioning Scale", ) guidance_scale = gr.Slider( minimum=0.0, maximum=25.0, step=0.25, value=7, label="Guidance Scale", ) seed = gr.Number( minimum=-1, maximum=9999999999, step=1, value=2313123, label="Seed", randomize=True, ) controlnet_start_0 = gr.Slider( minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="Controlnet Start 0", ) controlnet_start_1 = gr.Slider( minimum=0.0, maximum=1.0, step=0.01, value=0.0, label="Controlnet Start 1", ) controlnet_end_0 = gr.Slider( minimum=0.0, maximum=1.0, step=0.01, value=0.95, label="Controlnet End 0", ) controlnet_end_1 = gr.Slider( minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="Controlnet End 1", ) background = gr.Dropdown( label="Background", choices=backgrounds, value="gray", ) error_correction = gr.Dropdown( label="Error Correction", choices=correction_map.keys(), value="H", ) margin = gr.Slider( minimum=0, maximum=10, step=1, value=1, label="Margin", ) module_size = gr.Slider( minimum=1, maximum=100, step=1, value=16, label="Module Size", ) width = gr.Slider( minimum=512, maximum=1024, step=256, value=512, label="Width", ) height = gr.Slider( minimum=512, maximum=1024, step=256, value=512, label="Height", ) with gr.Row(): run_btn = gr.Button("Run") with gr.Column(): result_image = gr.Image(label="Result Image", elem_id="result_image") run_btn.click( inference_map, inputs=[ qr_code_content, prompt, negative_prompt, guidance_scale, controlnet_conditioning_scale, seed, controlnet_start_0, controlnet_start_1, controlnet_end_0, controlnet_end_1, background, error_correction, margin, module_size, width, height, ], outputs=[result_image], ) # login = os.environ.get("LOGIN", "admin") # password = os.environ.get("PASSWORD", "1234") blocks.queue(concurrency_count=1, max_size=40) blocks.launch(share=False) # blocks.launch(share=False, auth=(login, password))