File size: 1,002 Bytes
7cda8c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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,
).launch()