from controlnet_aux import OpenposeDetector from diffusers import StableDiffusionControlNetPipeline, ControlNetModel from diffusers import UniPCMultistepScheduler import gradio as gr import torch # Constants low_threshold = 100 high_threshold = 200 # Models pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet") controlnet = ControlNetModel.from_pretrained( "lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16 ) pipe = StableDiffusionControlNetPipeline.from_pretrained( "andite/anything-v4.0", controlnet=controlnet, torch_dtype=torch.float16 ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) # This command loads the individual model components on GPU on-demand. So, we don't # need to explicitly call pipe.to("cuda"). pipe.enable_model_cpu_offload() # Generator seed, generator = torch.manual_seed(0) def get_pose(image): return pose_model(image) def generate_images(image, prompt): pose = get_pose(image) output = pipe( prompt, pose, generator=generator, num_images_per_prompt=3 ) return output.images gr.Interface( generate_images, inputs=[ gr.Image(type="pil"), gr.Textbox( label="Enter your prompt", max_lines=1, placeholder="best quality, extremely detailed", ), ], outputs=gr.Gallery().style(grid=[2], height="auto"), title="Generate controlled outputs with ControlNet and Stable Diffusion. ", description="This Space uses pose estimated lines as the additional conditioning.", examples=[["yoga1.jpeg", "best quality, extremely detailed"]], allow_flagging=False, ).launch(enable_queue=True)