Spaces:
Sleeping
Sleeping
Commit
·
90bece4
1
Parent(s):
8bb6a15
Update main.py
Browse files
main.py
CHANGED
@@ -1,42 +1,51 @@
|
|
1 |
-
|
|
|
2 |
import requests
|
3 |
-
|
4 |
-
from io import BytesIO
|
5 |
-
import getpass, os
|
6 |
import warnings
|
|
|
|
|
|
|
|
|
7 |
from stability_sdk import client
|
8 |
import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
|
9 |
import replicate
|
10 |
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
upscale_engine="esrgan-v1-x2plus",
|
23 |
-
verbose=True,
|
24 |
-
)
|
25 |
-
|
26 |
-
# ClipDrop API function
|
27 |
-
def generate_image(prompt):
|
28 |
-
headers = {'x-api-key': api_key}
|
29 |
-
body_params = {'prompt': (None, prompt, 'text/plain')}
|
30 |
-
response = requests.post('https://clipdrop-api.co/text-to-image/v1', files=body_params, headers=headers)
|
31 |
-
|
32 |
-
if response.status_code == 200:
|
33 |
-
return Image.open(BytesIO(response.content))
|
34 |
-
else:
|
35 |
-
st.write(f"Request failed with status code {response.status_code}")
|
36 |
return None
|
37 |
|
38 |
-
|
39 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
40 |
answers = stability_api.upscale(init_image=img)
|
41 |
|
42 |
for resp in answers:
|
@@ -46,37 +55,20 @@ def upscale_image_stability(img):
|
|
46 |
"Your request activated the API's safety filters and could not be processed."
|
47 |
"Please submit a different image and try again.")
|
48 |
if artifact.type == generation.ARTIFACT_IMAGE:
|
49 |
-
|
|
|
|
|
|
|
|
|
50 |
|
51 |
-
|
52 |
-
def upscale_image_gfpgan(image_path):
|
53 |
-
with open(image_path, "rb") as img_file:
|
54 |
output = replicate.run(
|
55 |
"tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
|
56 |
input={"img": img_file, "version": "v1.4", "scale": 16}
|
57 |
)
|
58 |
-
response = requests.get(output)
|
59 |
-
return Image.open(BytesIO(response.content))
|
60 |
-
|
61 |
-
# Streamlit UI
|
62 |
-
st.title("Image Generator and Upscaler")
|
63 |
-
|
64 |
-
prompt = st.text_input("Enter a prompt for the image generation")
|
65 |
-
|
66 |
-
if st.button("Generate and Upscale"):
|
67 |
-
if prompt:
|
68 |
-
img1 = generate_image(prompt)
|
69 |
-
|
70 |
-
if img1:
|
71 |
-
st.image(img1, caption="Generated Image", use_column_width=True)
|
72 |
-
img1.save('generated_image.png')
|
73 |
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
|
78 |
-
|
79 |
-
st.image(img3, caption="Upscaled Image (GFPGAN)", use_column_width=True)
|
80 |
-
img3.save('upscaled_image_gfpgan.png')
|
81 |
-
else:
|
82 |
-
st.write("Please enter a prompt")
|
|
|
1 |
+
# main.py
|
2 |
+
|
3 |
import requests
|
4 |
+
import os
|
|
|
|
|
5 |
import warnings
|
6 |
+
import io
|
7 |
+
import getpass
|
8 |
+
|
9 |
+
from PIL import Image
|
10 |
from stability_sdk import client
|
11 |
import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
|
12 |
import replicate
|
13 |
|
14 |
+
def generate_and_upscale_image(text_prompt, clipdrop_api_key, stability_api_key, replicate_api_token):
|
15 |
+
|
16 |
+
headers = {'x-api-key': clipdrop_api_key}
|
17 |
+
body_params = {'prompt': (None, text_prompt, 'text/plain')}
|
18 |
+
|
19 |
+
response = requests.post('https://clipdrop-api.co/text-to-image/v1',
|
20 |
+
files=body_params,
|
21 |
+
headers=headers)
|
22 |
+
|
23 |
+
if response.status_code != 200:
|
24 |
+
print(f"Request failed with status code {response.status_code}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
25 |
return None
|
26 |
|
27 |
+
with open('generated_image.png', 'wb') as f:
|
28 |
+
f.write(response.content)
|
29 |
+
|
30 |
+
os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'
|
31 |
+
os.environ['STABILITY_KEY'] = stability_api_key
|
32 |
+
|
33 |
+
stability_api = client.StabilityInference(
|
34 |
+
key=os.environ['STABILITY_KEY'],
|
35 |
+
upscale_engine="esrgan-v1-x2plus",
|
36 |
+
verbose=True,
|
37 |
+
)
|
38 |
+
|
39 |
+
max_pixels = 1048576
|
40 |
+
img = Image.open('generated_image.png')
|
41 |
+
width, height = img.size
|
42 |
+
|
43 |
+
if width * height > max_pixels:
|
44 |
+
scale_factor = (max_pixels / (width * height))**0.5
|
45 |
+
new_width = int(width * scale_factor)
|
46 |
+
new_height = int(height * scale_factor)
|
47 |
+
img = img.resize((new_width, new_height))
|
48 |
+
|
49 |
answers = stability_api.upscale(init_image=img)
|
50 |
|
51 |
for resp in answers:
|
|
|
55 |
"Your request activated the API's safety filters and could not be processed."
|
56 |
"Please submit a different image and try again.")
|
57 |
if artifact.type == generation.ARTIFACT_IMAGE:
|
58 |
+
upscaled_img = Image.open(io.BytesIO(artifact.binary))
|
59 |
+
upscaled_img.save("upscaled_image.png")
|
60 |
+
|
61 |
+
os.environ['REPLICATE_API_TOKEN'] = replicate_api_token
|
62 |
+
Image.MAX_IMAGE_PIXELS = None
|
63 |
|
64 |
+
with open("upscaled_image.png", "rb") as img_file:
|
|
|
|
|
65 |
output = replicate.run(
|
66 |
"tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
|
67 |
input={"img": img_file, "version": "v1.4", "scale": 16}
|
68 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
69 |
|
70 |
+
response = requests.get(output)
|
71 |
+
final_img = Image.open(io.BytesIO(response.content))
|
72 |
+
final_img.save("gfpgan_upscaled_image.png")
|
73 |
|
74 |
+
return final_img
|
|
|
|
|
|
|
|