ms1449 commited on
Commit
c420f36
·
verified ·
1 Parent(s): f249c51

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +31 -26
app.py CHANGED
@@ -1,33 +1,39 @@
 
1
  import gradio as gr
2
- from diffusers import StableDiffusionPipeline
3
- import torch
4
  import os
5
- from huggingface_hub import login
6
 
7
- # Login using your Hugging Face token
8
- hf_token = os.getenv("HF_TOKEN") # This should be set in your Hugging Face Space Secrets
9
- login(token=hf_token)
10
 
11
- # Access the NVIDIA API key from the environment
12
- nvidia_api_key = os.getenv("NVIDIA_API_KEY")
 
 
13
 
14
- # Load the model using the API key and move it to GPU
15
- def load_model():
16
- model_id = "NVIDIA/sdxl-turbo"
17
- pipe = StableDiffusionPipeline.from_pretrained(
18
- model_id,
19
- torch_dtype=torch.float16,
20
- use_auth_token=nvidia_api_key # Use the API key here
21
- )
22
- pipe = pipe.to("cuda")
23
- return pipe
24
-
25
- pipe = load_model()
26
-
27
- # Function to generate a Kindle cover
28
  def generate_kindle_cover(prompt):
29
- image = pipe(prompt).images[0]
30
- return image
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
31
 
32
  # Create the Gradio interface
33
  iface = gr.Interface(
@@ -39,5 +45,4 @@ iface = gr.Interface(
39
  )
40
 
41
  # Launch the Gradio interface
42
- if __name__ == "__main__":
43
- iface.launch()
 
1
+ import requests
2
  import gradio as gr
3
+ from PIL import Image
4
+ from io import BytesIO
5
  import os
 
6
 
7
+ # Replace with your NVIDIA API key
8
+ API_KEY = os.getenv("NVIDIA_API_KEY") # Ensure the key is set in Space secrets
9
+ invoke_url = "https://ai.api.nvidia.com/v1/genai/stabilityai/sdxl-turbo"
10
 
11
+ headers = {
12
+ "Authorization": f"Bearer {API_KEY}",
13
+ "Accept": "application/json",
14
+ }
15
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
16
  def generate_kindle_cover(prompt):
17
+ payload = {
18
+ "text_prompts": [{"text": prompt}],
19
+ "seed": 0,
20
+ "sampler": "K_EULER_ANCESTRAL",
21
+ "steps": 2
22
+ }
23
+
24
+ response = requests.post(invoke_url, headers=headers, json=payload)
25
+
26
+ if response.status_code == 200:
27
+ response_body = response.json()
28
+ image_url = response_body.get('image_url') # Adjust based on actual response structure
29
+ if image_url:
30
+ image_response = requests.get(image_url)
31
+ image = Image.open(BytesIO(image_response.content))
32
+ return image
33
+ else:
34
+ return "No image URL found in response."
35
+ else:
36
+ return f"Error: {response.status_code} - {response.text}"
37
 
38
  # Create the Gradio interface
39
  iface = gr.Interface(
 
45
  )
46
 
47
  # Launch the Gradio interface
48
+ iface.launch()