Arshad Mangi commited on
Commit
8219a52
·
verified ·
1 Parent(s): 2747d31

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -0
app.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from diffusers import StableDiffusionPipeline
3
+ import torch
4
+ import os
5
+
6
+ # Retrieve Hugging Face API key from environment variables (Hugging Face Secrets)
7
+ hf_token = os.getenv("HF_TOKEN")
8
+
9
+ # Load the Stable Diffusion model in CPU mode
10
+ pipe = StableDiffusionPipeline.from_pretrained(
11
+ "stabilityai/stable-diffusion-2",
12
+ torch_dtype=torch.float32, # Ensure compatibility with CPU
13
+ use_auth_token=hf_token
14
+ )
15
+
16
+ # Force model to run on CPU
17
+ pipe.to("cpu")
18
+
19
+ # Optimize performance for CPU execution
20
+ pipe.enable_attention_slicing()
21
+
22
+ # Define image generation function
23
+ def generate_image(prompt):
24
+ try:
25
+ image = pipe(prompt).images[0]
26
+ return image
27
+ except Exception as e:
28
+ return f"Error: {str(e)}"
29
+
30
+ # Gradio Interface
31
+ iface = gr.Interface(
32
+ fn=generate_image,
33
+ inputs=gr.Textbox(label="Enter your prompt"),
34
+ outputs=gr.Image(label="Generated Image"),
35
+ title="Stable Diffusion (CPU Optimized)",
36
+ description="Generate AI-generated images using Stable Diffusion on CPU. No GPU required!"
37
+ )
38
+
39
+ if __name__ == "__main__":
40
+ iface.launch(share=True)