from diffusers import StableDiffusionXLPipeline
import gradio as gr
import torch
#from transformers.utils.hub import move_cache
import pkg_resources
#move_cache()

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 # Your prompt here
  neg_prompt = negative_prompt # Negative prompt here
  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()