Spaces:
Running
Running
| 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) |