File size: 2,846 Bytes
87ca7ca 7a279ff fd24b8d 87ca7ca fd24b8d 87ca7ca fd24b8d 87ca7ca fd24b8d 87ca7ca fd24b8d 87ca7ca fd24b8d 87ca7ca fd24b8d 87ca7ca 7a279ff 87ca7ca 7a279ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | import torch
from diffusers import StableDiffusionPipeline
import gradio as gr
# ---------------------------------------------------------------------------
# 1. Load the Stable Diffusion model from Hugging Face
# - We specify "runwayml/stable-diffusion-v1-5" as an example.
# - Use "revision='fp16'" and "torch_dtype=torch.float16" to use the half-precision weights.
# - .to('cuda') if GPU is available, else .to('cpu').
# ---------------------------------------------------------------------------
try:
pipe = StableDiffusionPipeline.from_pretrained(
"eric707/jibjab",
revision="fp16",
torch_dtype=torch.float16
).to("cuda")
device = "cuda"
except:
# If CUDA is not available, fall back to CPU (VERY slow for SD, but works in a pinch).
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16"
# If you're on CPU, you might remove the torch_dtype for better compatibility:
# torch_dtype=torch.float16 -> Not recommended on CPU
).to("cpu")
device = "cpu"
# ---------------------------------------------------------------------------
# 2. Define a function to generate images given a prompt.
# - We'll keep things simple and only accept a single prompt string.
# - Feel free to modify the inference steps, guidance scale, image size, etc.
# ---------------------------------------------------------------------------
def generate_image(prompt):
# Lower the inference steps or guidance scale if you run out of memory
image = pipe(
prompt,
num_inference_steps=30,
guidance_scale=7.5
).images[0]
return image
# ---------------------------------------------------------------------------
# 3. Build the Gradio UI
# - We use a Textbox for user input,
# and an Image component for displaying the generated image.
# ---------------------------------------------------------------------------
with gr.Blocks() as demo:
gr.Markdown("## Stable Diffusion Image Generation")
with gr.Row():
with gr.Column():
prompt_input = gr.Textbox(
label="Enter a prompt to generate an image",
placeholder="A photo of an astronaut riding a horse on Mars"
)
generate_button = gr.Button("Generate Image")
with gr.Column():
output_image = gr.Image(label="Generated Image")
generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image)
# ---------------------------------------------------------------------------
# 4. Launch the Gradio app
# ---------------------------------------------------------------------------
if __name__ == "__main__":
# By default, .launch() will pick up the PORT from the environment if on HF Spaces
demo.launch()
|