Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import requests | |
| from io import BytesIO | |
| from PIL import Image | |
| import os | |
| import tempfile | |
| import mimetypes | |
| TOKEN = os.getenv("TOKEN0") | |
| API_URL = os.getenv("API_URL") | |
| token_id = 0 | |
| tokens_tried = 0 | |
| no_of_accounts = 11 | |
| model_id = os.getenv("MODEL_ID") | |
| def get_image_from_url(url): | |
| """ | |
| Fetches and returns an image from a given URL, converting to PNG if needed. | |
| """ | |
| try: | |
| response = requests.get(url, stream=True) | |
| response.raise_for_status() | |
| image = Image.open(BytesIO(response.content)) | |
| return image, url # Return the image and the URL | |
| except requests.exceptions.RequestException as e: | |
| return f"Error fetching image: {e}", None | |
| except Exception as e: | |
| return f"Error processing image: {e}", None | |
| def generate_image(prompt, aspect_ratio, realism): | |
| global token_id | |
| global TOKEN | |
| global tokens_tried | |
| global no_of_accounts | |
| global model_id | |
| payload = { | |
| "id": model_id, | |
| "inputs": [prompt, aspect_ratio, str(realism).lower()], | |
| } | |
| headers = {"Authorization": f"Bearer {TOKEN}"} | |
| try: | |
| response_data = requests.post(API_URL, json=payload, headers=headers).json() | |
| if "error" in response_data: | |
| if 'error 429' in response_data['error']: | |
| if tokens_tried < no_of_accounts: | |
| token_id = (token_id + 1) % (no_of_accounts) | |
| tokens_tried += 1 | |
| TOKEN = os.getenv(f"TOKEN{token_id}") | |
| response_data = generate_image(prompt, aspect_ratio, realism) | |
| tokens_tried = 0 | |
| return response_data | |
| return "No credits available", None | |
| return response_data, None | |
| elif "output" in response_data: | |
| url = response_data['output'] | |
| image, url = get_image_from_url(url) | |
| return image, url # Return the image and the URL | |
| else: | |
| return "Error: Unexpected response from server", None | |
| except Exception as e: | |
| return f"Error", None | |
| def download_image(image_url): | |
| if not image_url: | |
| return None # Return None if image_url is empty | |
| try: | |
| response = requests.get(image_url, stream=True) | |
| response.raise_for_status() | |
| # Get the content type from the headers | |
| content_type = response.headers.get('content-type') | |
| extension = mimetypes.guess_extension(content_type) | |
| if not extension: | |
| extension = ".png" # Default to .png if can't determine the extension | |
| # Create a temporary file with the correct extension | |
| with tempfile.NamedTemporaryFile(suffix=extension, delete=False) as tmp_file: | |
| for chunk in response.iter_content(chunk_size=8192): | |
| tmp_file.write(chunk) | |
| temp_file_path = tmp_file.name | |
| return temp_file_path | |
| except Exception as e: | |
| return None | |
| # Define the Gradio interface | |
| interface = gr.Interface( | |
| fn=generate_image, | |
| inputs=[ | |
| gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate"), | |
| gr.Radio( | |
| choices=["1:1", "3:4", "4:3", "9:16", "16:9", "9:21", "21:9"], | |
| label="Aspect Ratio", | |
| value="16:9" # Default value | |
| ), | |
| gr.Checkbox(label="Realism", value=False), # Checkbox for realism (True/False) | |
| ], | |
| outputs=[ | |
| gr.Image(type="pil"), # Output image | |
| gr.File(label="Download Image", file_types = [".png", ".jpg", ".jpeg"]), # Output File to be downloaded | |
| ], | |
| title="Image Generator", | |
| description="Provide a prompt, select an aspect ratio, and set realism to generate an image.", | |
| ) | |
| # Launch the interface | |
| interface.launch() |