| import gradio as gr |
| from diffusers import DiffusionPipeline |
| import torch |
|
|
| model_id = "stabilityai/stable-diffusion-2-1-base" |
|
|
| pipe = DiffusionPipeline.from_pretrained( |
| model_id, |
| torch_dtype=torch.float16 |
| ) |
|
|
| if torch.cuda.is_available(): |
| pipe = pipe.to("cuda") |
| else: |
| pipe = pipe.to("cpu") |
|
|
|
|
| def generate_image(prompt): |
| if not prompt.strip(): |
| return None |
|
|
| image = pipe(prompt).images[0] |
| return image |
|
|
|
|
| with gr.Blocks(title="AI Image Generator") as demo: |
|
|
| gr.Markdown( |
| """ |
| # 🎨 AI Image Generator |
| Enter a prompt and let AI create an image for you! |
| """ |
| ) |
|
|
| prompt = gr.Textbox( |
| label="✨ Enter your prompt", |
| placeholder="A cute cottage in a magical forest at sunset..." |
| ) |
|
|
| generate_btn = gr.Button("🎨 Generate Image") |
|
|
| output = gr.Image(label="Generated Image") |
|
|
| generate_btn.click( |
| fn=generate_image, |
| inputs=prompt, |
| outputs=output |
| ) |
|
|
|
|
| demo.launch() |
