Singularity666/VisionGPT-Automation

1# main.py2 3import requests4import os5import warnings6import io7 8from PIL import Image9from stability_sdk import client10import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation11import replicate12 13def generate_and_upscale_image(text_prompt, clipdrop_api_key, stability_api_key, replicate_api_token):14 15    headers = {'x-api-key': clipdrop_api_key}16    body_params = {'prompt': (None, text_prompt, 'text/plain')}17 18    response = requests.post('https://clipdrop-api.co/text-to-image/v1',19                             files=body_params,20                             headers=headers)21 22    if response.status_code != 200:23        print(f"Request failed with status code {response.status_code}")24        return None, f"Request failed with status code {response.status_code}"25 26    with open('generated_image.png', 'wb') as f:27        f.write(response.content)28 29    os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'30    os.environ['STABILITY_KEY'] = stability_api_key31 32    stability_api = client.StabilityInference(33        key=os.environ['STABILITY_KEY'],34        upscale_engine="esrgan-v1-x2plus",35        verbose=True,36    )37 38    max_pixels = 104857639    img = Image.open('generated_image.png')40    width, height = img.size41 42    if width * height > max_pixels:43        scale_factor = (max_pixels / (width * height))**0.544        new_width = int(width * scale_factor)45        new_height = int(height * scale_factor)46        img = img.resize((new_width, new_height))47 48    answers = stability_api.upscale(init_image=img)49 50    for resp in answers:51        for artifact in resp.artifacts:52            if artifact.finish_reason == generation.FILTER:53                warnings.warn(54                    "Your request activated the API's safety filters and could not be processed."55                    "Please submit a different image and try again.")56            if artifact.type == generation.ARTIFACT_IMAGE:57                upscaled_img = Image.open(io.BytesIO(artifact.binary))58                upscaled_img.save("upscaled_image.png")59 60    os.environ['REPLICATE_API_TOKEN'] = replicate_api_token61    Image.MAX_IMAGE_PIXELS = None62 63    with open("upscaled_image.png", "rb") as img_file:64        output = replicate.run(65            "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",66            input={"img": img_file, "version": "v1.4", "scale": 16}67        )68 69    response = requests.get(output)70    if response.status_code != 200:71        return None, f"Failed to fetch upscaled image. Status code {response.status_code}"72 73    final_img = Image.open(io.BytesIO(response.content))74    img_byte_arr = io.BytesIO()75    final_img.save(img_byte_arr, format='PNG')76    img_byte_arr = img_byte_arr.getvalue()77 78    return img_byte_arr, None79