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()