Spaces:
Running
Running
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() | |