Spaces:
Runtime error
Runtime error
File size: 1,014 Bytes
b7aacca |
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 |
import gradio as gr
from PIL import Image
from authtoken import auth_token
import torch
import torch.cuda.amp as amp
from diffusers import StableDiffusionPipeline
modelid = "CompVis/stable-diffusion-v1-4"
device = torch.device("cpu") # Default to CPU device
if torch.cuda.is_available():
device = torch.device("cuda")
pipe = StableDiffusionPipeline.from_pretrained(modelid, use_auth_token=auth_token)
pipe.to(device)
def generate(prompt):
with torch.no_grad(), amp.autocast(enabled=device != torch.device("cpu")):
image = pipe(prompt, guidance_scale=8.5)["sample"][0]
image.save('generatedimage.png')
return image
def predict_text(prompt):
image = generate(prompt)
return image
def predict_image(input_image):
input_image.save('input_image.png')
prompt = input("Enter your prompt: ")
image = generate(prompt)
return image
iface = gr.Interface(
fn=predict_text,
inputs="text",
outputs="image",
capture_session=True,
)
iface.launch(share=True)
|