!pip install transformers==4.29.0
!pip install diffusers==0.16.1
!pip install accelerate==0.17.1
!pip install gradio==3.32.0
!pip install translate==3.6.1

import os
import moxing as mox
pretrained_models = ['mlsd', 'sd-controlnet-mlsd', 'stable-diffusion-v1-5']
for pretrained_model in pretrained_models:
 model_dir = os.path.join(os.getcwd(), pretrained_model)
 if not os.path.exists(model_dir):
 mox.file.copy_parallel(f'obs://modelarts-labs-bj4-v2/case_zoo/ControlNet/ControlNet_models/{pretrained_model}', model_dir)
 if os.path.exists(model_dir):
 print(f'{pretrained_model} download success')
 else:
 raise Exception(f'{pretrained_model} download Failed')
 else:
 print(f"{pretrained_model} already exists!")

import torch
from PIL import Image
from mlsd import MLSDdetector
from translate import Translator
from diffusers.utils import load_image
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
mlsd = MLSDdetector()
controlnet = ControlNetModel.from_pretrained("sd-controlnet-mlsd", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetPipeline.from_pretrained("stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()

from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
ori = Image.open("1920245540.jpg")
mlsd = MLSDdetector()
mlsd_img = mlsd(ori, thr_v=0.1, thr_d=0.1, detect_resolution=512, image_resolution=512, return_pil=True)
trans = Translator(from_lang="ZH",to_lang="EN-US")
prompt = "现代极简风卧室，床是黄色的，墙是浅咖色"
en_prompt = trans.translate(prompt)
gen_img = pipe(en_prompt, mlsd_img, num_inference_steps=20).images[0]
fig = plt.figure(figsize=(25, 10))
ax1 = fig.add_subplot(1, 3, 1)
plt.title('Orignial image', fontsize=16)
ax1.axis('off')
ax1.imshow(ori)
ax2 = fig.add_subplot(1, 3, 2)
plt.title('ML Detector image', fontsize=16)
ax2.axis('off')
ax2.imshow(mlsd_img)
ax3 = fig.add_subplot(1, 3, 3)
plt.title('Generate image', fontsize=16)
ax3.axis('off')
ax3.imshow(gen_img)
plt.show()

import gradio as gr
def mlsd(img, prompt, num_inference_steps, thr_v, thr_d, n_prompt, detect_resolution, image_resolution):
    trans = Translator(from_lang="ZH",to_lang="EN-US")
    prompt = trans.translate(prompt)
 n_prompt = trans.translate(n_prompt)
 mlsd = MLSDdetector()
 mlsd_img = mlsd(img, thr_v=0.1, thr_d=0.1, detect_resolution=512, image_resolution=512, return_pil=True)
 gen_img = pipe(prompt, mlsd_img, num_inference_steps=20, negative_prompt=n_prompt).images[0]
    result = [mlsd_img, gen_img]
 return result 
block = gr.Blocks().queue()
with block:
 with gr.Row():
 gr.Markdown("##  Control Stable Diffusion with MLSD")
 with gr.Row():
 with gr.Column():
 input_image = gr.Image(source='upload', type="numpy")
            prompt = gr.Textbox(label="描述")
 run_button = gr.Button(label="Run")
 with gr.Accordion("高级选项", open=False):
 num_inference_steps = gr.Slider(label="图像生成步数", minimum=1, maximum=100, value=20, step=1)
 value_threshold = gr.Slider(label="Hough value threshold (MLSD)", minimum=0.01, maximum=2.0, value=0.1, step=0.01)
 distance_threshold = gr.Slider(label="Hough distance threshold (MLSD)", minimum=0.01, maximum=20.0, value=0.1, step=0.01)
 n_prompt = gr.Textbox(label="否定提示",
                                      value='lowres, extra digit, fewer digits, cropped, worst quality, low quality')
 detect_resolution = gr.Slider(label="Hough Resolution", minimum=128, maximum=1024, value=512, step=1)
 image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512, step=64)
 with gr.Column():
 result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(columns=2, height='auto')
 ips = [input_image, prompt, num_inference_steps, value_threshold, distance_threshold, n_prompt, detect_resolution, image_resolution]
 run_button.click(fn=mlsd, inputs=ips, outputs=[result_gallery])
block.launch(share=True)