app.py

import os
import random

import cv2
import gradio as gr
import numpy as np
import torch
# import spaces #[uncomment to use ZeroGPU]
from diffusers import (ControlNetModel, StableDiffusionControlNetPipeline,
                       StableDiffusionPipeline)
from peft import PeftModel
from PIL import Image

device = "cuda" if torch.cuda.is_available() else "cpu"
# model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
# model_repo_id = "CompVis/stable-diffusion-v1-4"
# model_dropdown = ["stabilityai/sdxl-turbo", "CompVis/stable-diffusion-v1-4"]

models = [
    # "gstranger/kawaiicat-lora-1.4",
    "CompVis/stable-diffusion-v1-4",
    "stabilityai/sdxl-turbo",
    "sd-legacy/stable-diffusion-v1-5",
]
controlnet_modes = ["canny", "Line Art"]

model_dropdown = [
    "stabilityai/sdxl-turbo",
    "CompVis/stable-diffusion-v1-4",
    "sd-legacy/stable-diffusion-v1-5",
]

def process_control_image(image, mode="canny"):
  if mode == "canny":
    image = np.array(image)
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    canny = cv2.Canny(blurred, 50, 150)
    return Image.fromarray(canny)
  return image

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

# pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
# pipe = pipe.to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024
MODEL_NAME = "CompVis/stable-diffusion-v1-4"
CKPT_DIR = "sd-14-lora-1000"



# @spaces.GPU #[uncomment to use ZeroGPU]
def infer(
    model_id,
    prompt,
    negative_prompt,
    randomize_seed=False,
    width=512,
    height=512,
    seed = 488,
    guidance_scale=7,
    num_inference_steps=50,

    lora_enable=True,
    lora_scale=0.8,
    
    controlnet_enable=False,
    control_mode="Line Art",
    control_strength=0.8,
    control_image=None,

    ip_adapter_enable=False,
    ip_adapter_scale=0.8,
    ip_image=None,


    torch_dtype=torch_dtype,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    else:
        seed = 488

    generator = torch.Generator().manual_seed(seed)

    params = {'prompt': prompt,
            'negative_prompt': negative_prompt,
            'guidance_scale': guidance_scale,
            'num_inference_steps': num_inference_steps,
            'width': width,
            'height': height,
            'generator': generator,
            }
    
    print("in infer 1")
    print("controlnet_enable", controlnet_enable)

    controlnet = None
    if controlnet_enable and control_image is not None:
        print("in controlnet_enable")
        if control_mode == "canny":
            controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_canny", 
                                                         torch_dtype=torch_dtype, cache_dir="./models_cache")
            
        else: 
            control_mode == "Line Art"
            controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_lineart", 
                                                         torch_dtype=torch_dtype, cache_dir="./models_cache")
            
        pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id, 
                                                                 controlnet=controlnet,
                                                                 torch_dtype=torch_dtype, 
                                                                 safety_checker=None) #.to(device)
        
        params['image'] = process_control_image(control_image, control_mode)
        params['controlnet_conditioning_scale'] = float(control_strength)


    else:
        print("step: basic pipeline")
        pipe = StableDiffusionPipeline.from_pretrained(model_id, 
                                                       torch_dtype=torch_dtype, 
                                                       safety_checker=None) #.to(device)
        print("step: basic pipeline done!")

    if lora_enable:
        print("step: lora")
        unet_sub_dir = os.path.join(CKPT_DIR, "unet")
        text_encoder_sub_dir = os.path.join(CKPT_DIR, "text_encoder")
        adapter_name="sd-14-lora"

        pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
        pipe.text_encoder = PeftModel.from_pretrained(
            pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name
        )
        params['cross_attention_kwargs']={"scale": lora_scale}
        print("step: lora done!")

    if torch_dtype in (torch.float16, torch.bfloat16):
        pipe.unet.half()
        pipe.text_encoder.half()

    if ip_adapter_enable:
        print("step: ip_adapter_enable")
        pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter-plus_sd15.bin")
        pipe.set_ip_adapter_scale(ip_adapter_scale)
        params['ip_adapter_image'] = process_control_image(ip_image, "")
        print("step: ip_adapter_enable done!")


    # pipe.to(device)

    print("step: start generating")
    print(params)
    image = pipe(**params
    ).images[0]
    print("step: generating done!")
    return image, seed


examples = [
    "kawaiicat. The cat is sitting. The cat's tail is curled up at the end. The cat is pleased and is enjoying its time.",
    "kawaiicat. The cat is sitting upright. The cat is eating some noodles with the chopsticks from a green bowl, which it's holding in his hands.",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # Text-to-Image kawaiicat Stickers")
        with gr.Row():
            # Dropdown to select the model from Hugging Face
            model_id = gr.Dropdown(
                label="Model",
                choices=models,
                value=models[0],  # Default model
            )

        lora_scale = gr.Slider(
            label="LORA Scale",
            minimum=0,
            maximum=1,
            step=0.01,
            value=0.8,
        )

        lora_enable = gr.Checkbox(label="Use LORA", value=True)

        with gr.Column():
            controlnet_enable = gr.Checkbox(label="Enable ControlNet", value=False)
            with gr.Accordion("ControlNet Settings", visible=False) as controlnet_accordion:
                control_mode = gr.Dropdown(controlnet_modes, label="Control Mode", value="canny")
                control_strength = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="Control Strength")
                control_image = gr.Image(label="Control Image", type="pil")

            ip_adapter_enable = gr.Checkbox(label="Enable IP-Adapter", value=False)
            with gr.Accordion("IP-Adapter Settings", visible=False) as ipadapter_accordion:
                ip_adapter_scale = gr.Slider(0, 1, value=0.5, label="IP-Adapter Scale")
                ip_image = gr.Image(label="Reference Image", type="pil")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
                value="kawaiicat. The cat is having fun, is smiling."
            )        
            
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter your negative prompt",
                value="bad anatomy, crop image, bad face of the cat"
            )


            run_button = gr.Button("Run", scale=0, variant="primary")

        result = gr.Image(label="Result", show_label=False)

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
                visible=True,
                value="bad anatomy, crop image, bad face of the cat"
            )

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=42,
            )

            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)

            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=512,  # Replace with defaults that work for your model
                )

                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=512,  # Replace with defaults that work for your model
                )

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=10.0,  # Replace with defaults that work for your model
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=50,  # Replace with defaults that work for your model
                )

        gr.Examples(examples=examples, inputs=[prompt])

        controlnet_enable.change(
            lambda x: gr.update(visible=x), 
            controlnet_enable, 
            controlnet_accordion
        )
        ip_adapter_enable.change(
            lambda x: gr.update(visible=x), 
            ip_adapter_enable, 
            ipadapter_accordion
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            model_id,
            prompt,
            negative_prompt,
            randomize_seed,
            width,
            height,
            seed,
            guidance_scale,
            num_inference_steps,
            lora_enable, lora_scale,
            controlnet_enable, control_mode, control_strength, control_image,
            ip_adapter_enable, ip_adapter_scale, ip_image

        ],
        outputs=[result, seed],
    )


if __name__ == "__main__":
    demo.launch()