Shriharsh commited on
Commit
b7aacca
1 Parent(s): 4a8685f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +39 -0
app.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from PIL import Image
3
+ from authtoken import auth_token
4
+ import torch
5
+ import torch.cuda.amp as amp
6
+ from diffusers import StableDiffusionPipeline
7
+
8
+ modelid = "CompVis/stable-diffusion-v1-4"
9
+ device = torch.device("cpu") # Default to CPU device
10
+ if torch.cuda.is_available():
11
+ device = torch.device("cuda")
12
+ pipe = StableDiffusionPipeline.from_pretrained(modelid, use_auth_token=auth_token)
13
+ pipe.to(device)
14
+
15
+ def generate(prompt):
16
+ with torch.no_grad(), amp.autocast(enabled=device != torch.device("cpu")):
17
+ image = pipe(prompt, guidance_scale=8.5)["sample"][0]
18
+
19
+ image.save('generatedimage.png')
20
+ return image
21
+
22
+ def predict_text(prompt):
23
+ image = generate(prompt)
24
+ return image
25
+
26
+ def predict_image(input_image):
27
+ input_image.save('input_image.png')
28
+ prompt = input("Enter your prompt: ")
29
+ image = generate(prompt)
30
+ return image
31
+
32
+ iface = gr.Interface(
33
+ fn=predict_text,
34
+ inputs="text",
35
+ outputs="image",
36
+ capture_session=True,
37
+ )
38
+
39
+ iface.launch(share=True)