hsuwill000's picture
Update app.py
86d4979 verified
import gradio as gr
from gradio_client import Client
from PIL import Image
import os
import time
import traceback
import asyncio
from huggingface_hub import HfApi
# Your Hugging Face API key (ensure this is set in your environment or replace directly)
api_key = os.getenv('MY_API_KEY')
api = HfApi(token=api_key)
# List of repos (private spaces)
repos = [
"hsuwill000/Fluently-v4-LCM-openvino_0",
"hsuwill000/Fluently-v4-LCM-openvino_1",
"hsuwill000/Fluently-v4-LCM-openvino_2",
"hsuwill000/Fluently-v4-LCM-openvino_3",
"hsuwill000/Fluently-v4-LCM-openvino_4",
]
class CustomClient(Client):
def __init__(self, *args, timeout=30, **kwargs):
super().__init__(*args, **kwargs)
self.timeout = timeout
def _request(self, method, url, **kwargs):
kwargs['timeout'] = self.timeout
return super()._request(method, url, **kwargs)
# Counter for image filenames to avoid overwriting
count = 0
async def infer_single_gradio(client, prompt):
global count
# Prepare the inputs for the prediction
inputs = {
"prompt": prompt,
#"num_inference_steps": 10 # Number of inference steps for the model
}
try:
# Send the request to the model and receive the image
result = await asyncio.to_thread(client.predict, inputs, api_name="/infer")
# Open the resulting image
image = Image.open(result)
# Create a unique filename to save the image
filename = f"img_{count:08d}.jpg"
while os.path.exists(filename):
count += 1
filename = f"img_{count:08d}.jpg"
# Save the image locally
image.save(filename)
print(f"Saved image as {filename}")
# Return the image to be displayed in Gradio
return image
except Exception as e:
# Handle any errors that occur
print(f"An exception occurred: {str(e)}")
print("Stack trace:")
traceback.print_exc() # Print stack trace for debugging
return None # Return nothing if an error occurs
async def infer_gradio(prompt: str, output_images: gr.Gallery):
#can't work Clear previous images from the gallery
#output_images.update([])
# Record the start time
start_time = time.time()
# Create a list of tasks (one for each repo)
tasks = []
for repo in repos:
runtime_info = api.get_space_runtime(repo)
if runtime_info.stage == 'APP_STARTING':
continue
if runtime_info.stage in ['SLEEPING', 'PAUSED']:
print(f"{repo} is now SLEEPING or PAUSED.")
api.restart_space(repo_id=repo)
continue
# Create a CustomClient instance for each repo
client = CustomClient(repo, hf_token=api_key, timeout=300)
task = infer_single_gradio(client, prompt)
tasks.append(task)
# Run all tasks concurrently (i.e., generate images from all repos)
results = await asyncio.gather(*tasks)
# Calculate the time taken for image generation
end_time = time.time()
time_taken = end_time - start_time
print(f"Time taken to generate the image(s): {time_taken:.2f} seconds")
# Return the results (images)
return results # Return all the images as a list
# Define Gradio Interface
with gr.Blocks() as demo:
with gr.Row(): # Use a Row to place the prompt input and the button side by side
prompt_input = gr.Textbox(
label="Enter Your Prompt",
show_label="False",
placeholder="Type your prompt for image generation here",
lines=1, # Set the input to be only one line tall
interactive=True # Allow user to interact with the textbox
)
# Change the button text to "RUN:" and align it with the prompt input
run_button = gr.Button("RUN")
# Output image display area (will show multiple images)
output_images = gr.Gallery(label="Generated Images", elem_id="gallery", show_label=False)
# Connecting the button click to the image generation function
run_button.click(infer_gradio, inputs=[prompt_input], outputs=output_images)
# Launch Gradio app
demo.launch()