|
from diffusers import StableDiffusionXLPipeline |
|
import gradio as gr |
|
import torch |
|
|
|
|
|
import pkg_resources |
|
|
|
|
|
print("------------------------") |
|
package_name = "transformers" |
|
version = pkg_resources.get_distribution(package_name).version |
|
print(f"The version of {package_name} is: {version}") |
|
print("------------------------") |
|
|
|
def segMindImage(prompt, negative_prompt): |
|
pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B", torch_dtype=torch.float16, use_safetensors=True, variant="fp16") |
|
pipe.to("cuda") |
|
prompt = prompt |
|
neg_prompt = negative_prompt |
|
image = pipe(prompt=prompt, negative_prompt=neg_prompt).images[0] |
|
|
|
return image |
|
|
|
app = gr.Interface(segMindImage, |
|
inputs = [gr.Text(label="Prompt",placeholder="Write Prompt"),gr.Text(label="Negative Prompt",placeholder="Write Negative Prompt")], |
|
outputs = gr.Image(label="Image"), |
|
css =".gradio-container {background-image: linear-gradient(#7F7FD5, #91EAE4, #A3C9E2)}", |
|
theme = gr.themes.Soft(), |
|
title = "SD Image Diffusion") |
|
|
|
app.launch() |
|
|