import gradio as gr
from diffusers import AutoPipelineForText2Image
from diffusers.schedulers import DPMSolverMultistepScheduler
from generate_propmts import generate_prompt  # Assuming you have this module
from PIL import Image
import asyncio
import threading
import traceback

# Define the SchedulerWrapper class
class SchedulerWrapper:
    def __init__(self, scheduler):
        self.scheduler = scheduler
        self._step = threading.local()
        self._step.step = 0

    def __getattr__(self, name):
        return getattr(self.scheduler, name)

    def step(self, *args, **kwargs):
        try:
            self._step.step += 1
            return self.scheduler.step(*args, **kwargs)
        except IndexError:
            self._step.step = 0
            return self.scheduler.step(*args, **kwargs)

    @property
    def timesteps(self):
        return self.scheduler.timesteps

    def set_timesteps(self, *args, **kwargs):
        return self.scheduler.set_timesteps(*args, **kwargs)

# Load the model and wrap the scheduler
model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")

scheduler = DPMSolverMultistepScheduler.from_config(model.scheduler.config)
wrapped_scheduler = SchedulerWrapper(scheduler)
model.scheduler = wrapped_scheduler

# Define the image generation function
async def generate_image(prompt):
    try:
        num_inference_steps = 5  # Adjust this value as needed

        # Use the model to generate an image
        output = await asyncio.to_thread(
            model,
            prompt=prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=0.0,  # Typical value for guidance scale in image generation
            output_type="pil"  # Directly get PIL Image objects
        )

        # Check for output validity and return
        if output.images:
            return output.images[0]
        else:
            raise Exception("No images returned by the model.")
    except Exception as e:
        print(f"Error generating image: {e}")
        traceback.print_exc()
        return None  # Return None on error to handle it gracefully in the UI

# Define the inference function
async def inference(sentence_mapping, character_dict, selected_style):
    images = []
    print(f'sentence_mapping: {sentence_mapping}, character_dict: {character_dict}, selected_style: {selected_style}')
    prompts = []

    # Generate prompts for each paragraph
    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(prompt)
        print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")

    # Use asyncio.gather to run generate_image in parallel
    tasks = [generate_image(prompt) for prompt in prompts]
    images = await asyncio.gather(*tasks)

    # Filter out None values
    images = [image for image in images if image is not None]

    return images

# Define the Gradio interface
gradio_interface = gr.Interface(
    fn=inference,
    inputs=[
        gr.JSON(label="Sentence Mapping"),
        gr.JSON(label="Character Dict"),
        gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
    ],
    outputs=gr.Gallery(label="Generated Images")
)

# Run the Gradio app
if __name__ == "__main__":
    gradio_interface.launch()