import gradio as gr
import torch
from diffusers import AutoencoderKL, StableDiffusionXLPipeline,StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_single_file(
    "https://huggingface.co/ethe/Architecture_model/blob/main/architectureExterior_v40Exterior.safetensors",
    # "/mnt/pfs-guan-ssai/cv/panxuhao/checkpoints/stable-diffusion-xl-base-1.0/sd_xl_base_1.0.safetensors",
    torch_dtype=torch.float32,
    # local_files_only=True,
    variant="fp16", 
    use_safetensors=True,
)

num_images_per_prompt = 4
positive_prompt = 'Modern Elegance:Explore the seamless blend of sleek lines, minimalist aesthetics, and cutting-edge design in modern interior decor'

negative_prompt = '(nsfw:1.3),(Nude:1.3),(Naked:1.3),low quality, blurry, bad anatomy, worst quality, text, watermark, normal quality, ugly, signature, lowres, deformed, disfigured, cropped, jpeg artifacts, error, \
                    mutation, logo, watermark,text, logo,contact, error, blurry, cropped, username, artist name, (worst quality, low quality:1.4),monochrome,'
def inference(prompt,ref_image):
    prompt  += positive_prompt
    image = pipe(
        prompt=prompt, 
        negative_prompt=negative_prompt,
        num_inference_steps=20, 
        # cross_attention_kwargs={"scale": lora_scale}, 
        generator=torch.manual_seed(0),
        width=512,
        height=512,
        guidance_scale=7.0,
        use_karras_sigmas=True,
        num_images_per_prompt=num_images_per_prompt,  # 如果是 4, image 就是四张图片组成的 List
    ).images
    return image

with gr.Row():
    input_prompt = gr.Textbox(placeholder="输入你对室内设计的要求，以便 AI 能够更好地满足你的需求。", label="要求")
with gr.Row():
    ref_image = gr.Image(height=512, width=512, label="参考图片")
    
    result_image = gr.Gallery(label="可能满足你需求的室内设计图：", 
                                    columns=3, 
                                    height="auto", 
                                    object_fit="contain")
demo = gr.Interface(
    fn=inference,
    inputs=[input_prompt,ref_image],
    outputs=[result_image]
)

demo.launch()