File size: 1,844 Bytes
4f72200 b3b7325 4f72200 b3b7325 c0fb55e b3b7325 c0fb55e b3b7325 c0fb55e b3b7325 c0fb55e b3b7325 c0fb55e b3b7325 |
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 49 50 51 52 53 54 55 56 57 58 59 |
import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch
def load_model():
try:
# Load the model without forcing half-precision
pipeline = StableDiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1",
torch_dtype=torch.float32, # Use float32 for CPU compatibility
safety_checker=None # Disable safety checker for faster inference
)
except Exception as e:
print(f"Error loading the model: {e}")
raise
# Configure the scheduler for faster generation
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
# Move to CPU
try:
pipeline = pipeline.to("cpu")
except Exception as e:
print(f"Error moving the model to device: {e}")
raise
return pipeline
# Initialize the model
try:
model = load_model()
except Exception as e:
print(f"Error initializing the model: {e}")
# Define Gradio interface
def generate(prompt, guidance_scale=7.5, num_inference_steps=50):
try:
# Generate image from the prompt
images = model(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images
return images[0]
except Exception as e:
return f"Error generating image: {e}"
# 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()
|