Singularity666 commited on
Commit
90bece4
·
1 Parent(s): 8bb6a15

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +50 -58
main.py CHANGED
@@ -1,42 +1,51 @@
1
- import streamlit as st
 
2
  import requests
3
- from PIL import Image
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
- # API keys
12
- api_key = 'YOUR_API_KEY'
13
- os.environ['STABILITY_KEY'] = 'YOUR_API_KEY'
14
- os.environ['REPLICATE_API_TOKEN'] = 'REPLICATE_API_TOKEN' # Replace with your actual API token
15
-
16
- # Increase the pixel limit
17
- Image.MAX_IMAGE_PIXELS = None
18
-
19
- # Establish connection to Stability API
20
- stability_api = client.StabilityInference(
21
- key=os.environ['STABILITY_KEY'],
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
- # Stability API function
39
- def upscale_image_stability(img):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
- return Image.open(io.BytesIO(artifact.binary))
 
 
 
 
50
 
51
- # GFPGAN function
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
- img2 = upscale_image_stability(img1)
75
- st.image(img2, caption="Upscaled Image (Stability API)", use_column_width=True)
76
- img2.save('upscaled_image_stability.png')
77
 
78
- img3 = upscale_image_gfpgan('upscaled_image_stability.png')
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