File size: 1,565 Bytes
9175f45 b9a6225 9175f45 669fd03 b9a6225 669fd03 b9a6225 3aaaf76 b9a6225 0a24053 b9a6225 669fd03 0a24053 b9a6225 0a24053 b9a6225 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch
def load_model():
# Specify the Stable Diffusion pipeline with an appropriate model type
pipeline = StableDiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1",
torch_dtype=torch.float16,
revision="fp16",
safety_checker=None # Disable safety checker if necessary
)
# Set the scheduler (optional but recommended)
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
# Move pipeline to GPU or ZeroGPU
pipeline = pipeline.to("cuda") # or ZeroGPU-specific setup
return pipeline
# Initialize the model
try:
model = load_model()
except Exception as e:
print(f"Error loading the model: {e}")
# Define Gradio interface
def generate(prompt, guidance_scale=7.5, num_inference_steps=50):
# Generate the image
images = model(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images
return images[0]
# Gradio Interface
with gr.Blocks() as demo:
with gr.Row():
prompt = gr.Textbox(label="Enter your prompt")
guidance_scale = gr.Slider(1.0, 10.0, value=7.5, label="Guidance Scale")
steps = gr.Slider(10, 100, value=50, label="Number of Inference Steps")
with gr.Row():
submit = gr.Button("Generate")
with gr.Row():
output = gr.Image()
submit.click(generate, inputs=[prompt, guidance_scale, steps], outputs=output)
demo.launch()
|