| 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() | |