import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load the stable diffusion model
pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", 
    revision="fp16", 
    torch_dtype=torch.float16
)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Create a Gradio interface
gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=2, placeholder="Enter an exercise like 'child to cobra pose'"),
    outputs="image",
    title="AI Exercise Visual Generator",
    description="Enter an exercise name. This tool will generate a visual representation using AI."
).launch()