Spaces:
Runtime error
Runtime error
| import os | |
| import asyncio | |
| import time | |
| from generate_prompts import generate_prompt | |
| from diffusers import AutoPipelineForText2Image | |
| from io import BytesIO | |
| import gradio as gr | |
| import ray | |
| ray.init() | |
| class ModelActor: | |
| def __init__(self): | |
| """ | |
| Initializes the ModelActor class and loads the text-to-image model. | |
| """ | |
| self.model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo") | |
| async def generate_image(self, prompt, prompt_name): | |
| """ | |
| Generates an image based on the provided prompt. | |
| Parameters: | |
| - prompt (str): The input text for image generation. | |
| - prompt_name (str): A name for the prompt, used for logging. | |
| Returns: | |
| bytes: The generated image data in bytes format, or None if generation fails. | |
| """ | |
| start_time = time.time() | |
| process_id = os.getpid() | |
| try: | |
| output = await self.model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0) | |
| if isinstance(output.images, list) and len(output.images) > 0: | |
| image = output.images[0] | |
| buffered = BytesIO() | |
| image.save(buffered, format="JPEG") | |
| image_bytes = buffered.getvalue() | |
| end_time = time.time() | |
| return image_bytes | |
| else: | |
| return None | |
| except Exception as e: | |
| return None | |
| async def queue_api_calls(sentence_mapping, character_dict, selected_style): | |
| """ | |
| Generates images for all provided prompts in parallel using Ray actors. | |
| Parameters: | |
| - sentence_mapping (dict): Mapping between paragraph numbers and sentences. | |
| - character_dict (dict): Dictionary mapping characters to their descriptions. | |
| - selected_style (str): Selected illustration style. | |
| Returns: | |
| dict: A dictionary where keys are paragraph numbers and values are image data in bytes format. | |
| """ | |
| prompts = [] | |
| for paragraph_number, sentences in sentence_mapping.items(): | |
| combined_sentence = " ".join(sentences) | |
| prompt = generate_prompt(combined_sentence, sentence_mapping, character_dict, selected_style) | |
| prompts.append((paragraph_number, prompt)) | |
| num_prompts = len(prompts) | |
| num_actors = min(num_prompts, 20) | |
| model_actors = [ModelActor.remote() for _ in range(num_actors)] | |
| tasks = [model_actors[i % num_actors].generate_image.remote(prompt, f"Prompt {paragraph_number}") for i, (paragraph_number, prompt) in enumerate(prompts)] | |
| responses = await asyncio.gather(*[asyncio.to_thread(ray.get, task) for task in tasks]) | |
| images = {paragraph_number: response for (paragraph_number, _), response in zip(prompts, responses)} | |
| return images | |
| def process_prompt(sentence_mapping, character_dict, selected_style): | |
| """ | |
| Processes the provided prompts and generates images. | |
| Parameters: | |
| - sentence_mapping (dict): Mapping between paragraph numbers and sentences. | |
| - character_dict (dict): Dictionary mapping characters to their descriptions. | |
| - selected_style (str): Selected illustration style. | |
| Returns: | |
| dict: A dictionary where keys are paragraph numbers and values are image data in bytes format. | |
| """ | |
| try: | |
| loop = asyncio.get_running_loop() | |
| except RuntimeError: | |
| loop = asyncio.new_event_loop() | |
| asyncio.set_event_loop(loop) | |
| cmpt_return = loop.run_until_complete(queue_api_calls(sentence_mapping, character_dict, selected_style)) | |
| return cmpt_return | |
| gradio_interface = gr.Interface( | |
| fn=process_prompt, | |
| inputs=[gr.JSON(label="Sentence Mapping"), gr.JSON(label="Character Dict"), gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")], | |
| outputs="json" | |
| ) | |
| if __name__ == "__main__": | |
| gradio_interface.launch() | |