text-to-image / app.py
sudo-soldier's picture
Update app.py
e1c2d0e verified
raw
history blame
947 Bytes
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4-original", torch_dtype=torch.float32)
pipe.to("cpu") # Make sure to use CPU
# Define the function for inference
def infer(prompt, guidance_scale=7.5, num_inference_steps=50):
# Generate image from the prompt
image = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images[0]
return image
# Create Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("# Hyper-Sketch with Stable Diffusion")
with gr.Row():
prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=2)
generate_button = gr.Button("Generate Image")
output_image = gr.Image(label="Generated Image")
generate_button.click(infer, inputs=[prompt], outputs=[output_image])
# Launch the app
demo.launch()