Spaces:
Runtime error
Runtime error
| 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", | |
| "n0madic/deliberate", | |
| #"SG161222/Realistic_Vision_V1.4", | |
| 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, | |
| controlnet_conditioning_scale_1: 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, controlnet_conditioning_scale_1), | |
| 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_0 = gr.Slider( | |
| minimum=0.5, | |
| maximum=2.5, | |
| step=0.01, | |
| value=1.5, | |
| label="Controlnet Conditioning Scale", | |
| ) | |
| controlnet_conditioning_scale_1 = gr.Slider( | |
| minimum=0.5, | |
| maximum=2.5, | |
| step=0.01, | |
| value=1.0, | |
| label="Controlnet Conditioning Scale (corners)", | |
| ) | |
| 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_0, | |
| controlnet_conditioning_scale_1, | |
| 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)) | |