Spaces:
Runtime error
Runtime error
Commit
·
23659de
1
Parent(s):
f5140e7
Create main.py
Browse files
main.py
ADDED
@@ -0,0 +1,104 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
import os
|
3 |
+
import requests
|
4 |
+
from PIL import Image
|
5 |
+
from io import BytesIO
|
6 |
+
|
7 |
+
# Set up environment variables for API keys
|
8 |
+
os.environ['CLIPDROP_API_KEY'] = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
|
9 |
+
os.environ['STABILITY_API_KEY'] = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
|
10 |
+
os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'
|
11 |
+
|
12 |
+
# Importing Replicate and Stability SDK libraries
|
13 |
+
import replicate
|
14 |
+
import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
|
15 |
+
|
16 |
+
def generate_image(prompt):
|
17 |
+
# Make a POST request to the ClipDrop text-to-image API
|
18 |
+
url = 'https://clipdrop-api.co/text-to-image/v1'
|
19 |
+
headers = {'x-api-key': os.environ['CLIPDROP_API_KEY']}
|
20 |
+
data = {'prompt': prompt}
|
21 |
+
response = requests.post(url, headers=headers, data=data)
|
22 |
+
|
23 |
+
if response.status_code == 200:
|
24 |
+
# Get the generated image from the response
|
25 |
+
img = Image.open(BytesIO(response.content))
|
26 |
+
return img
|
27 |
+
else:
|
28 |
+
st.error('Failed to generate image from text prompt.')
|
29 |
+
|
30 |
+
def upscale_image(img):
|
31 |
+
# Set up Stability API connection
|
32 |
+
stability_api = replicate.StabilityInference(
|
33 |
+
key=os.environ['STABILITY_API_KEY'],
|
34 |
+
upscale_engine="stable-diffusion-x4-latent-upscaler"
|
35 |
+
)
|
36 |
+
|
37 |
+
# Upscale the image using the Stability API
|
38 |
+
upscale_responses = stability_api.upscale(init_image=img)
|
39 |
+
|
40 |
+
if upscale_responses:
|
41 |
+
# Get the upscaled image from the response
|
42 |
+
upscaled_img = None
|
43 |
+
for resp in upscale_responses:
|
44 |
+
for artifact in resp.artifacts:
|
45 |
+
if artifact.type == generation.ARTIFACT_IMAGE:
|
46 |
+
upscaled_img = Image.open(BytesIO(artifact.binary))
|
47 |
+
break
|
48 |
+
if upscaled_img:
|
49 |
+
break
|
50 |
+
return upscaled_img
|
51 |
+
else:
|
52 |
+
st.error('Failed to upscale the image.')
|
53 |
+
|
54 |
+
def gfpgan_upscale_image(img):
|
55 |
+
# Set up Replicate API connection
|
56 |
+
replicate_api = replicate.ReplicateAPI(token=os.environ['REPLICATE_API_TOKEN'])
|
57 |
+
|
58 |
+
# Upscale the image using GFPGAN model
|
59 |
+
output = replicate_api.run(
|
60 |
+
"tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
|
61 |
+
input={"img": img, "version": "v1.4", "scale": 16}
|
62 |
+
)
|
63 |
+
|
64 |
+
if output:
|
65 |
+
# Get the upscaled image URI from the output
|
66 |
+
response = requests.get(output)
|
67 |
+
upscaled_img = Image.open(BytesIO(response.content))
|
68 |
+
return upscaled_img
|
69 |
+
else:
|
70 |
+
st.error('Failed to upscale the image using GFPGAN.')
|
71 |
+
|
72 |
+
def main():
|
73 |
+
st.title("Image Upscaling")
|
74 |
+
st.write("Enter a text prompt to generate and upscale an image.")
|
75 |
+
|
76 |
+
# Get user input for the text prompt
|
77 |
+
prompt = st.text_input("Enter a text prompt:", max_chars=1000)
|
78 |
+
|
79 |
+
if st.button("Generate and Upscale"):
|
80 |
+
# Generate the image from text prompt using ClipDrop API
|
81 |
+
generated_img = generate_image(prompt)
|
82 |
+
|
83 |
+
if generated_img:
|
84 |
+
st.image(generated_img, caption='Generated Image', use_column_width=True)
|
85 |
+
|
86 |
+
# Upscale the generated image using Stability API
|
87 |
+
upscaled_img = upscale_image(generated_img)
|
88 |
+
|
89 |
+
if upscaled_img:
|
90 |
+
st.image(upscaled_img, caption='Upscaled Image (Stability API)', use_column_width=True)
|
91 |
+
|
92 |
+
# Further upscale the image using GFPGAN
|
93 |
+
gfpgan_upscaled_img = gfpgan_upscale_image(upscaled_img)
|
94 |
+
|
95 |
+
if gfpgan_upscaled_img:
|
96 |
+
st.image(gfpgan_upscaled_img, caption='Upscaled Image (GFPGAN)', use_column_width=True)
|
97 |
+
st.success("Image generation and upscaling completed successfully.")
|
98 |
+
else:
|
99 |
+
st.error("Failed to further upscale the image using GFPGAN.")
|
100 |
+
else:
|
101 |
+
st.error("Failed to upscale the generated image.")
|
102 |
+
|
103 |
+
if __name__ == "__main__":
|
104 |
+
main()
|