Singularity666 commited on
Commit
23659de
·
1 Parent(s): f5140e7

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +104 -0
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()