ahabb commited on
Commit
5a4f96a
·
verified ·
1 Parent(s): 09dde59

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +70 -0
app.py ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import requests
3
+ import base64
4
+ from PIL import Image
5
+ from io import BytesIO
6
+ from dotenv import load_dotenv
7
+ import os
8
+
9
+ # Load the .env file
10
+ load_dotenv()
11
+
12
+ # Get the API key from the environment variable
13
+ api_key = os.getenv("API_KEY")
14
+
15
+ # Define the endpoint URL
16
+ invoke_url = "https://ai.api.nvidia.com/v1/genai/stabilityai/sdxl-turbo"
17
+
18
+ # Define the function that interacts with the API
19
+ def generate_image(prompt, seed=0, sampler="K_EULER_ANCESTRAL", steps=2):
20
+ # Define the headers with the correct API key
21
+ headers = {
22
+ "Authorization": f"Bearer {api_key}",
23
+ "Accept": "application/json",
24
+ }
25
+
26
+ # Define the payload for the POST request
27
+ payload = {
28
+ "text_prompts": [{"text": prompt}],
29
+ "seed": seed,
30
+ "sampler": sampler,
31
+ "steps": steps
32
+ }
33
+
34
+ # Make the POST request to the API
35
+ response = requests.post(invoke_url, headers=headers, json=payload)
36
+
37
+ # Raise an error if the request was unsuccessful
38
+ response.raise_for_status()
39
+
40
+ # Get the response body
41
+ response_body = response.json()
42
+
43
+ # Extract the base64 string from the response
44
+ base64_str = response_body['artifacts'][0]['base64']
45
+
46
+ # Decode the base64 string
47
+ image_data = base64.b64decode(base64_str)
48
+
49
+ # Open the image using PIL
50
+ image = Image.open(BytesIO(image_data))
51
+
52
+ return image
53
+
54
+ # Define the Gradio interface
55
+ inputs = [
56
+ gr.Textbox(label="Prompt", placeholder="Enter your image description here"),
57
+ gr.Number(label="Seed", value=0, precision=0),
58
+ gr.Dropdown(label="Sampler", choices=["K_EULER_ANCESTRAL", "DDIM", "PLMS"], value="K_EULER_ANCESTRAL"),
59
+ gr.Slider(label="Steps", minimum=1, maximum=100, value=2)
60
+ ]
61
+
62
+ outputs = gr.Image(label="Generated Image")
63
+
64
+ gr.Interface(
65
+ fn=generate_image,
66
+ inputs=inputs,
67
+ outputs=outputs,
68
+ title="Stable Diffusion XL Image Generator",
69
+ description="Generate images using Stable Diffusion XL by providing a text prompt."
70
+ ).launch()