Spaces:
Sleeping
Sleeping
import gradio as gr | |
from diffusers import StableDiffusionPipeline | |
import torch | |
# Initialize the Stable Diffusion pipeline | |
model_id = "sd-legacy/stable-diffusion-v1-5" | |
# Check if CUDA is available | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
print(f"Using device: {device}") | |
# Initialize pipeline with appropriate settings based on device | |
if device == "cuda": | |
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16) | |
else: | |
pipe = StableDiffusionPipeline.from_pretrained(model_id) | |
print("Running on CPU - using float32 precision") | |
pipe = pipe.to(device) | |
def generate_image(prompt, negative_prompt="", num_inference_steps=50, seed=None, guidance_scale=7.5): | |
""" | |
Generate an image from a text prompt using Stable Diffusion | |
""" | |
try: | |
# Set the seed if provided | |
if seed is not None and seed.strip(): | |
generator = torch.Generator(device).manual_seed(int(seed)) | |
else: | |
generator = None | |
# Generate the image | |
image = pipe( | |
prompt, | |
negative_prompt=negative_prompt, | |
num_inference_steps=num_inference_steps, | |
guidance_scale=guidance_scale, | |
generator=generator | |
).images[0] | |
return image | |
except Exception as e: | |
print(f"Error generating image: {str(e)}") | |
return None | |
# Create the Gradio interface | |
with gr.Blocks(title="Stable Diffusion Text-to-Image") as demo: | |
gr.Markdown("# Stable Diffusion Text-to-Image Generator") | |
gr.Markdown(f"Running on: {device.upper()}") | |
gr.Markdown("Enter a text prompt to generate an image using Stable Diffusion") | |
with gr.Row(): | |
with gr.Column(): | |
prompt = gr.Textbox( | |
label="Enter your prompt", | |
placeholder="A beautiful sunset over mountains, digital art", | |
lines=3 | |
) | |
neg_prompt = gr.Textbox( | |
label="Enter your negative prompt", | |
value="ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame,extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature,cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face", | |
lines=3, | |
visible=False | |
) | |
steps = gr.Slider( | |
minimum=2, | |
maximum=50, | |
value=20, | |
step=1, | |
label="Number of inference steps" | |
) | |
seed = gr.Textbox( | |
label="Seed (optional)", | |
placeholder="Leave empty for random seed" | |
) | |
generate_btn = gr.Button("Generate Image", variant='primary') | |
with gr.Column(): | |
output = gr.Image(label="Generated Image") | |
generate_btn.click( | |
fn=generate_image, | |
inputs=[prompt, neg_prompt, steps, seed], | |
outputs=output | |
) | |
if __name__ == "__main__": | |
demo.launch(share=True) |