Singularity666 commited on
Commit
33c27ec
1 Parent(s): 27fea1c

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +47 -0
main.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import replicate
3
+ import os
4
+ import requests
5
+ from PIL import Image
6
+ from io import BytesIO
7
+
8
+ # Set up environment variables for Replicate and Stability AI API Tokens
9
+ os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I' # Replace with your actual Replicate API token
10
+ os.environ['STABILITY_KEY'] = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO' # Replace with your actual Stability AI API key
11
+
12
+ def upscale_image(image_path):
13
+ # Open the image file
14
+ with open(image_path, "rb") as img_file:
15
+ # Upscale the image using Stability AI's API and Stable Diffusion 4x Upscaler
16
+ stability_api_url = "https://api.stability.ai/v1/image-upscaling"
17
+ headers = {
18
+ 'Authorization': f'Bearer {os.environ["STABILITY_KEY"]}',
19
+ }
20
+ files = {
21
+ 'image': img_file,
22
+ 'upscaler': 'stable-diffusion-x4-latent-upscaler',
23
+ }
24
+ response = requests.post(stability_api_url, headers=headers, files=files)
25
+
26
+ # Retrieve the upscaled image from the response
27
+ img = Image.open(BytesIO(response.content))
28
+ img.save("upscaled.png") # Save the upscaled image
29
+ return img
30
+
31
+ def main():
32
+ st.title("Image Upscaling")
33
+ st.write("Upload an image and enter a prompt to generate and upscale an image.")
34
+
35
+ uploaded_file = st.file_uploader("Choose an image...", type="png")
36
+ prompt = st.text_input("Enter a prompt for image generation")
37
+
38
+ if uploaded_file is not None and prompt:
39
+ with open("temp_img.png", "wb") as f:
40
+ f.write(uploaded_file.getbuffer())
41
+ st.success("Uploaded image successfully!")
42
+ if st.button("Generate and Upscale Image"):
43
+ img = upscale_image("temp_img.png")
44
+ st.image(img, caption='Generated and Upscaled Image', use_column_width=True)
45
+
46
+ if __name__ == "__main__":
47
+ main()