Spaces:
Sleeping
Sleeping
import gradio as gr | |
from huggingface_hub import login | |
import os | |
hf_token = os.environ.get("HF_TOKEN") | |
login(token=hf_token) | |
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL | |
from diffusers.utils import load_image | |
from PIL import Image | |
import torch | |
import numpy as np | |
import cv2 | |
controlnet = ControlNetModel.from_pretrained( | |
"diffusers/controlnet-canny-sdxl-1.0", | |
torch_dtype=torch.float16 | |
) | |
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) | |
pipe = StableDiffusionXLControlNetPipeline.from_pretrained( | |
"stabilityai/stable-diffusion-xl-base-1.0", | |
controlnet=controlnet, | |
vae=vae, | |
torch_dtype=torch.float16, | |
) | |
custom_model = "fffiloni/eugene_jour_general" | |
# This is where you load your trained weights | |
pipe.load_lora_weights(custom_model, use_auth_token=True) | |
#pipe.to("cuda") | |
pipe.enable_model_cpu_offload() | |
def infer(image_in, prompt): | |
prompt = prompt | |
negative_prompt = "" | |
image = load_image("https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png") | |
controlnet_conditioning_scale = 0.5 # recommended for good generalization | |
image = np.array(image) | |
image = cv2.Canny(image, 100, 200) | |
image = image[:, :, None] | |
image = np.concatenate([image, image, image], axis=2) | |
image = Image.fromarray(image) | |
images = pipe( | |
prompt, negative_prompt=negative_prompt, image=image, controlnet_conditioning_scale=controlnet_conditioning_scale, | |
).images | |
#images[0].save(f"hug_lab.png") | |
return images[0] | |
with gr.Blocks() as demo: | |
with gr.Column(): | |
image_in = gr.Image(source="upload", type="filepath") | |
prompt = gr.Textbox(label="Prompt") | |
submit_btn = gr.Button("Submit") | |
result = gr.Image(label="Result") | |
submit_btn.click( | |
fn = infer, | |
inputs = [image_in, prompt], | |
outputs = [result] | |
) | |
demo.queue().launch() | |