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.float16,
    # local_files_only=True,
    variant="fp16", 
    use_safetensors=True,
)

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

def inference():
    image = pipe(
        prompt=prompt, 
        negative_prompt="watermark, logo, symbol, word, phrase, human, noisy, blurry, deformed, ugly, NSFW, low quality, worst quality, monochrome, illustration, sketch, grayscale, backlight, messy, blurry",
        num_inference_steps=30, 
        # cross_attention_kwargs={"scale": lora_scale}, 
        generator=torch.manual_seed(0),
        width=768,
        height=768,
        guidance_scale=7.0,
        use_karras_sigmas=True,
        num_images_per_prompt=num_images_per_prompt,  # 如果是 4, image 就是四张图片组成的 List
    ).images
    return image
demo = gr.Interface(
    fn=inference,
    inputs=gr.Image(height=512, width=512, label="参考图片"),
    outputs=gr.Gallery(label="可能满足你需求的室内设计图：", 
                                  columns=3, 
                                  height="auto", 
                                  object_fit="contain")
)

demo.launch()