ms1449 commited on
Commit
a70b44b
·
verified ·
1 Parent(s): 814d990

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ # Load the model using the API key and move it to GPU
12
+ def load_model():
13
+ model_id = "NVIDIA/sdxl-turbo"
14
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_auth_token=True)
15
+ pipe = pipe.to("cuda")
16
+ return pipe
17
+
18
+ pipe = load_model()
19
+
20
+ # Function to generate a Kindle cover
21
+ def generate_kindle_cover(prompt):
22
+ image = pipe(prompt).images[0]
23
+ return image
24
+
25
+ # Create the Gradio interface
26
+ iface = gr.Interface(
27
+ fn=generate_kindle_cover,
28
+ inputs="text",
29
+ outputs="image",
30
+ title="Kindle Cover Generator",
31
+ description="Generate high-quality covers for Amazon Kindle books using the NVIDIA SDXL-Turbo model."
32
+ )
33
+
34
+ # Launch the Gradio interface
35
+ if __name__ == "__main__":
36
+ iface.launch()