Singularity666 commited on
Commit
bad15c0
·
1 Parent(s): 680ecae

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +7 -7
main.py CHANGED
@@ -10,6 +10,7 @@ import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
10
  # Configure your API keys here
11
  CLIPDROP_API_KEY = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
12
  STABLE_DIFFUSION_API_KEY = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
 
13
 
14
  # Set up environment variable for Replicate API Token
15
  os.environ['REPLICATE_API_TOKEN'] = 'r8_Tm3LQMS81QaGXzzdGVRyUCOQ3cuNd1i1sJlqp' # Replace with your actual API token
@@ -39,15 +40,14 @@ def resize_image(image_bytes, max_size=(256, 256)):
39
  img.save(buffer, format="PNG")
40
  return buffer.getvalue()
41
 
42
- def upscale_image_stable_diffusion(image_bytes):
43
  # Set up environment variables
44
- os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'
45
- os.environ['STABILITY_KEY'] = STABLE_DIFFUSION_API_KEY
46
 
47
  # Set up the connection to the API
48
  stability_api = client.StabilityInference(
49
- key=os.environ['STABILITY_KEY'],
50
- upscale_engine="stable-diffusion-x4-latent-upscaler",
51
  verbose=True,
52
  )
53
 
@@ -96,8 +96,8 @@ def main():
96
  st.success("Resizing image...")
97
  resized_image_bytes = resize_image(image_bytes)
98
 
99
- st.success("Upscaling image with stable-diffusion-x4-latent-upscaler...")
100
- upscaled_image_bytes = upscale_image_stable_diffusion(resized_image_bytes) # Change this line
101
 
102
  st.success("Further upscaling image with GFPGAN...")
103
  img = further_upscale_image(upscaled_image_bytes)
 
10
  # Configure your API keys here
11
  CLIPDROP_API_KEY = '1143a102dbe21628248d4bb992b391a49dc058c584181ea72e17c2ccd49be9ca69ccf4a2b97fc82c89ff1029578abbea'
12
  STABLE_DIFFUSION_API_KEY = 'sk-GBmsWR78MmCSAWGkkC1CFgWgE6GPgV00pNLJlxlyZWyT3QQO'
13
+ ESRGAN_API_KEY = 'your_esrgan_api_key_here'
14
 
15
  # Set up environment variable for Replicate API Token
16
  os.environ['REPLICATE_API_TOKEN'] = 'r8_Tm3LQMS81QaGXzzdGVRyUCOQ3cuNd1i1sJlqp' # Replace with your actual API token
 
40
  img.save(buffer, format="PNG")
41
  return buffer.getvalue()
42
 
43
+ def upscale_image_esrgan(image_bytes):
44
  # Set up environment variables
45
+ os.environ['ESRGAN_API_KEY'] = ESRGAN_API_KEY
 
46
 
47
  # Set up the connection to the API
48
  stability_api = client.StabilityInference(
49
+ key=os.environ['ESRGAN_API_KEY'],
50
+ upscale_engine="esrgan-v1-x2plus",
51
  verbose=True,
52
  )
53
 
 
96
  st.success("Resizing image...")
97
  resized_image_bytes = resize_image(image_bytes)
98
 
99
+ st.success("Upscaling image with ESRGAN...")
100
+ upscaled_image_bytes = upscale_image_esrgan(resized_image_bytes)
101
 
102
  st.success("Further upscaling image with GFPGAN...")
103
  img = further_upscale_image(upscaled_image_bytes)