Upload control_app.py

control_app.py

@@ -0,0 +1,131 @@
+'''
+!git clone https://huggingface.co/spaces/radames/SPIGA-face-alignment-headpose-estimator
+!cp -r SPIGA-face-alignment-headpose-estimator/SPIGA .
+!pip install -r SPIGA/requirements.txt
+!pip install datasets
+!huggingface-cli login
+'''
+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.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", "戴帽子穿灰色衣服的男子"]
+]
+
+from PIL import Image
+def pred_func(image, prompt):
+    out = single_pred_features(image)
+    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,
+)
+gr.launch(share=False)
+
+if __name__ == "__main__":
+    '''
+    control_image = load_image("./conditioning_image_1.png")
+    prompt = "戴眼镜的中年男子"
+    # generate image
+    generator = torch.manual_seed(0)
+    image = pipe(
+         prompt, num_inference_steps=50, generator=generator, image=control_image
+    ).images[0]
+    image
+
+    control_image = load_image("./conditioning_image_1.png")
+    prompt = "穿蓝色衣服的秃头男子"
+    # generate image
+    generator = torch.manual_seed(0)
+    image = pipe(
+         prompt, num_inference_steps=50, generator=generator, image=control_image
+    ).images[0]
+    image
+
+    control_image = load_image("./conditioning_image_2.png")
+    prompt = "金色头发的美丽女子"
+    # generate image
+    generator = torch.manual_seed(0)
+    image = pipe(
+         prompt, num_inference_steps=50, generator=generator, image=control_image
+    ).images[0]
+    image
+
+    control_image = load_image("./conditioning_image_2.png")
+    prompt = "绿色运动衫的男子"
+    # generate image
+    generator = torch.manual_seed(0)
+    image = pipe(
+         prompt, num_inference_steps=50, generator=generator, image=control_image
+    ).images[0]
+    image
+
+    from huggingface_hub import HfApi
+    hf_api = HfApi()
+
+    hf_api.upload_file(
+        path_or_fileobj = "TSD_save_only/diffusion_pytorch_model.bin",
+        path_in_repo = "diffusion_pytorch_model.bin",
+        repo_id = "svjack/ControlNet-Face-Zh",
+        repo_type = "model",
+    )
+
+    hf_api.upload_file(
+        path_or_fileobj = "TSD_save_only/config.json",
+        path_in_repo = "config.json",
+        repo_id = "svjack/ControlNet-Face-Zh",
+        repo_type = "model",
+    )
+    '''
+    pass