ControlNet / app.py
ayyuce's picture
Removed source parameter
1c649b3 verified
import gradio as gr
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
import torch
model_id = "runwayml/stable-diffusion-v1-5"
controlnet_id = "lllyasviel/control_v11p_sd15_openpose"
controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float32)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
model_id,
controlnet=controlnet,
safety_checker=None,
torch_dtype=torch.float32
)
pipe = pipe.to("cpu")
pipe.enable_attention_slicing()
def generate_image(prompt, control_image, num_inference_steps=25, guidance_scale=7.5, controlnet_conditioning_scale=1.0):
"""
Generate an image using the ControlNet pipeline.
Args:
prompt (str): Your text prompt for image generation.
control_image (PIL.Image): A control image to guide generation.
num_inference_steps (int): Number of denoising steps.
guidance_scale (float): Classifier-free guidance scale.
controlnet_conditioning_scale (float): How strongly to condition on the control image.
Returns:
PIL.Image: The generated image.
"""
result = pipe(
prompt=prompt,
image=control_image,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
controlnet_conditioning_scale=controlnet_conditioning_scale
)
return result.images[0]
with gr.Blocks() as demo:
gr.Markdown("# ControlNet Image Generator on CPU\nThis demo uses a ControlNet pipeline (openpose variant) with Stable Diffusion to generate images guided by a control image. Note: Running on CPU can be slow!")
with gr.Row():
prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your image prompt here", value="A futuristic cityscape at dusk")
with gr.Row():
control_image_input = gr.Image(label="Control Image", type="pil")
output_image = gr.Image(label="Generated Image", type="pil")
with gr.Row():
num_steps = gr.Slider(minimum=10, maximum=50, value=25, step=1, label="Inference Steps")
guidance = gr.Slider(minimum=1.0, maximum=15.0, value=7.5, step=0.5, label="Guidance Scale")
control_scale = gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="ControlNet Conditioning Scale")
generate_btn = gr.Button("Generate Image")
generate_btn.click(
fn=generate_image,
inputs=[prompt_input, control_image_input, num_steps, guidance, control_scale],
outputs=output_image
)
demo.launch()