File size: 732 Bytes
4cc7b30
 
 
 
 
 
 
788dc63
98e6a84
788dc63
 
4cc7b30
 
 
f4326d7
4cc7b30
 
 
 
 
 
 
 
 
 
 
 
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
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

model_name = "runwayml/stable-diffusion-v1-5"
device = "cuda" if torch.cuda.is_available() else "cpu"

if device == "cuda":
    pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
else:
    pipe = StableDiffusionPipeline.from_pretrained(model_name).to(device)

def generate_image(prompt):
    with torch.no_grad():
        image = pipe(prompt, num_inference_steps=50, guidance_scale=9.5).images[0]

    return image

iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="BL's T2I Generator",
    description="Enter a prompt to generate an image."
)

iface.launch()