Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from diffusers import StableDiffusionPipeline | |
| import torch | |
| from PIL import Image | |
| import os | |
| import uuid | |
| # Load Stable Diffusion v1.5 | |
| pipe = StableDiffusionPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", | |
| safety_checker=None, | |
| torch_dtype=torch.float32 | |
| ).to("cpu") | |
| # Output folder | |
| os.makedirs("outputs", exist_ok=True) | |
| # Aspect ratio mapping | |
| aspect_ratios = { | |
| "Square (512x512)": (512, 512), | |
| "Portrait (512x768)": (512, 768), | |
| "Landscape (768x512)": (768, 512) | |
| } | |
| def generate_image(user_prompt, ratio): | |
| width, height = aspect_ratios[ratio] | |
| # Enhance prompt internally for better quality | |
| enhanced_prompt = f"{user_prompt}, highly detailed, ultra realistic, cinematic lighting, sharp focus" | |
| image = pipe(prompt=enhanced_prompt, height=height, width=width, guidance_scale=7.5).images[0] | |
| # Save image | |
| filename = f"outputs/{uuid.uuid4().hex}.png" | |
| image.save(filename) | |
| return image | |
| # UI | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# ๐ Stable Diffusion v1.5 - Image Generator") | |
| with gr.Row(): | |
| prompt = gr.Textbox(label="Prompt", placeholder="A shining star in the sky") | |
| ratio = gr.Radio(list(aspect_ratios.keys()), label="Aspect Ratio", value="Square (512x512)") | |
| with gr.Row(): | |
| generate_btn = gr.Button("Generate Image") | |
| with gr.Column(): | |
| output_img = gr.Image(label="Generated Image", interactive=False) | |
| generate_btn.click(fn=generate_image, inputs=[prompt, ratio], outputs=output_img) | |
| demo.launch() | |