AP123 commited on
Commit
ad68f85
1 Parent(s): 83bbc9e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +49 -0
app.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from PIL import Image
4
+ from diffusers import DiffusionPipeline
5
+ import os
6
+
7
+ # Constants
8
+ #SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", "0") == "1"
9
+
10
+ # Initialize the model
11
+ pipe = DiffusionPipeline.from_pretrained(
12
+ "playgroundai/playground-v2.5-1024px-aesthetic",
13
+ torch_dtype=torch.float16,
14
+ variant="fp16",
15
+ ).to("cuda")
16
+
17
+ # Safety Checker (if necessary)
18
+ #if SAFETY_CHECKER:
19
+ # Implement or import the safety checker code here
20
+
21
+ def generate_image(prompt, num_inference_steps=50, guidance_scale=7):
22
+ # Generate image
23
+ results = pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)
24
+
25
+ # Safety check (if necessary)
26
+ if SAFETY_CHECKER:
27
+ # Implement the safety check logic here
28
+ pass
29
+
30
+ return results.images[0]
31
+
32
+ # Gradio Interface
33
+ with gr.Blocks() as demo:
34
+ gr.Markdown("## Playground-V2.5 Demo")
35
+ with gr.Row():
36
+ prompt = gr.Textbox(label='Enter your image prompt')
37
+ num_inference_steps = gr.Slider(minimum=1, maximum=75, step=1, label='Number of Inference Steps', value=50)
38
+ guidance_scale = gr.Slider(minimum=1, maximum=20, step=0.1, label='Guidance Scale', value=7)
39
+ submit = gr.Button('Generate Image')
40
+ img = gr.Image(label='Generated Image')
41
+
42
+ submit.click(
43
+ fn=generate_image,
44
+ inputs=[prompt, num_inference_steps, guidance_scale],
45
+ outputs=img,
46
+ )
47
+
48
+
49
+ demo.queue().launch()