from pred_color import * import gradio as gr from diffusers import ( AutoencoderKL, ControlNetModel, DDPMScheduler, StableDiffusionControlNetPipeline, UNet2DConditionModel, UniPCMultistepScheduler, ) import torch from diffusers.utils import load_image controlnet_model_name_or_path = "svjack/ControlNet-Face-Zh" controlnet = ControlNetModel.from_pretrained(controlnet_model_name_or_path) #controlnet = controlnet.to("cuda") base_model_path = "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1" pipe = StableDiffusionControlNetPipeline.from_pretrained( base_model_path, controlnet=controlnet, #torch_dtype=torch.float16 ) # speed up diffusion process with faster scheduler and memory optimization pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) pipe.safety_checker = None #pipe.enable_model_cpu_offload() #pipe = pipe.to("cuda") if torch.cuda.is_available(): pipe = pipe.to("cuda") else: #pipe.enable_model_cpu_offload() pass example_sample = [ #["Protector_Cromwell_style.png", "戴帽子穿灰色衣服的男子"] ["boy.jpeg", "戴帽子穿灰色衣服的男子"] ] from PIL import Image def pred_func(image, prompt): #out = single_pred_features(image) features ,face_features = single_pred_features(image) req_img = produce_center_crop_image(features ,face_features) out = {} out["spiga_seg"] = req_img if type(out) == type({}): #return out["spiga_seg"] control_image = out["spiga_seg"] if type(image) == type("") and os.path.exists(image): image = Image.open(image).convert("RGB") elif hasattr(image, "shape"): image = Image.fromarray(image).convert("RGB") else: image = image.convert("RGB") image = image.resize((512, 512)) generator = torch.manual_seed(0) image = pipe( prompt, num_inference_steps=50, generator=generator, image=control_image ).images[0] return control_image ,image gr=gr.Interface(fn=pred_func, inputs=['image','text'], outputs=[gr.Image(label='output').style(height=512), gr.Image(label='output').style(height=512)], examples=example_sample if example_sample else None, cache_examples = False ) gr.launch(share=False)