test_gradio / app.py
amos1088's picture
test gradio
52d3f89
raw
history blame
1.82 kB
import gradio as gr
import torch
from diffusers import StableDiffusion3Pipeline, ControlNetModel, UniPCMultistepScheduler
from huggingface_hub import login
import os
# Log in to Hugging Face with token from environment variables
token = os.getenv("HF_TOKEN")
login(token=token)
# Model IDs for the base Stable Diffusion model and ControlNet variant
model_id = "stabilityai/stable-diffusion-3.5-large-turbo"
controlnet_id = "lllyasviel/control_v11p_sd15_inpaint" # Adjust based on ControlNet needs
# Load ControlNet and Stable Diffusion models
controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.bfloat16)
pipe = StableDiffusion3Pipeline.from_pretrained(model_id, controlnet=controlnet, torch_dtype=torch.bfloat16)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda") if torch.cuda.is_available() else pipe
# Gradio interface function
def generate_image(prompt, reference_image):
# Prepare the reference image
reference_image = reference_image.convert("RGB").resize((512, 512))
# Generate the image using the pipeline with ControlNet
generated_image = pipe(
prompt=prompt,
image=reference_image,
controlnet_conditioning_scale=1.0,
guidance_scale=7.5,
num_inference_steps=50
).images[0]
return generated_image
# Set up Gradio interface
interface = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(label="Prompt"),
gr.Image(type="pil", label="Reference Image (Style)")
],
outputs="image",
title="Image Generation with Stable Diffusion 3.5 and ControlNet",
description="Generates an image based on a text prompt and style reference image using Stable Diffusion 3.5 and ControlNet."
)
# Launch the Gradio interface
interface.launch()