painter / app_old.py
waterplayfire's picture
Update app_old.py
a008657
import gradio as gr
import os
import torch
from transformers import pipeline
from diffusers import StableDiffusionPipeline
from gradio.components import Image
from gradio.components import Textbox
model_id = "CompVis/stable-diffusion-v1-4"
my_model = "waterplayfire/MyModel"
device = "cuda"
cuda_support = torch.cuda.is_available()
title = "Txt to Image"
pipeline = None
if cuda_support:
title = "Txt to Image by GPU"
pipeline = StableDiffusionPipeline.from_pretrained(my_model, torch_dtype=torch.float16, revision="fp16")
pipeline = pipeline.to(device)
torch.backends.cudnn.enabled = True
else:
title = "Txt to Image by CPU"
pipeline = StableDiffusionPipeline.from_pretrained(my_model)
#title = os.environ['model_fetch']
def predict(prompt):
predictions = pipeline(prompt, guidance_scale=7.5).images[0]
return predictions#{p["label"]: p["score"] for p in predictions}
gr.Interface(
predict,
inputs=Textbox(),
outputs=Image(),
title=title,
layout="vertical",
).launch()