app.py

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