import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Load the model (replace with your actual model path)
model_name = "alvdansen/midsommarcartoon"  # Using the specified model
device = "cpu"  # Ensure we are using CPU

print(f"Using device: {device}")

try:
    # Load the Stable Diffusion model
    print("Loading model...")
    pipe = StableDiffusionPipeline.from_pretrained(model_name).to(device)
    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")
    raise

def generate_image(prompt):
    try:
        print(f"Generating image for prompt: {prompt}")
        # Generate image based on the prompt
        with torch.no_grad():
            image = pipe(prompt).images[0]
        return image
    except Exception as e:
        print(f"Error generating image: {e}")
        return None

# Gradio interface setup
gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter your prompt (e.g., 'stickman in a field')"),
    outputs=gr.Image(type="pil"),
    title="Doodle Stickman Image Generator",
    description="Generate doodle-style stickman images from text prompts."
).launch()