import gradio as gr
import requests
import torch
from diffusers import StableDiffusionXLPipeline, DPMSolverSinglestepScheduler

# Constants
OLLAMA_HOST = "http://127.0.0.1:11434"
DEFAULT_MODEL = "your_default_model_name"

# Function to check Ollama server status
def check_server_status():
    try:
        response = requests.get(f"{OLLAMA_HOST}/api/tags")
        if response.status_code == 200:
            return "Ollama server is running smoothly."
        else:
            return f"Server error: {response.status_code}"
    except Exception as e:
        return f"Error connecting to server: {str(e)}"

# Function to handle chatbot conversation
def handle_conversation(user_input):
    try:
        # Construct API request to Ollama
        response = requests.post(
            f"{OLLAMA_HOST}/api/chat",
            json={
                "model": DEFAULT_MODEL,
                "messages": [
                    {"role": "system", "content": "Enhance the following prompt for image generation:"},
                    {"role": "user", "content": user_input}
                ]
            }
        )
        response.raise_for_status()
        reply = response.json().get('choices', [{}])[0].get('message', {}).get('content', '')
        return reply
    except Exception as e:
        return f"Error in chatbot response: {str(e)}"

# Load the Stable Diffusion model
pipe = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash", torch_dtype=torch.float16).to("cuda")
pipe.scheduler = DPMSolverSinglestepScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")

# Function to generate image based on prompt and advanced settings
def generate_image(prompt, steps, guidance_scale, width, height, seed, randomize_seed):
    try:
        generator = torch.manual_seed(seed) if not randomize_seed else None
        # Image generation
        image = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance_scale, width=width, height=height, generator=generator).images[0]
        image_path = "output.png"
        image.save(image_path)
        return image_path
    except Exception as e:
        return f"Error generating image: {str(e)}"

# Gradio interface
with gr.Blocks() as app:
    with gr.Row():
        input_box = gr.Textbox(label="Enter your prompt")
        run_button = gr.Button("Run")
        enhanced_prompt_output = gr.Textbox(label="Enhanced Prompt")
        image_output = gr.Image(label="Generated Image")

    def handle_request(user_input, steps, guidance_scale, width, height, seed, randomize_seed):
        enhanced_prompt = handle_conversation(user_input)
        if "Error" in enhanced_prompt:
            return enhanced_prompt, None
        image_path = generate_image(enhanced_prompt, steps, guidance_scale, width, height, seed, randomize_seed)
        return enhanced_prompt, image_path

    # Advanced settings
    with gr.Accordion("Advanced Settings", open=False):
        steps = gr.Slider(1, 50, value=7, step=1, label="Number of inference steps")
        guidance_scale = gr.Slider(1, 20, value=3.5, step=0.1, label="Guidance Scale")
        width = gr.Slider(64, 2048, value=1024, step=64, label="Width")
        height = gr.Slider(64, 2048, value=1024, step=64, label="Height")
        seed = gr.Number(value=0, label="Seed")
        randomize_seed = gr.Checkbox(value=True, label="Randomize seed")

    run_button.click(handle_request, inputs=[input_box, steps, guidance_scale, width, height, seed, randomize_seed], outputs=[enhanced_prompt_output, image_output])

    # Example prompts
    with gr.Row():
        gr.Examples(examples=[
            ["a tiny astronaut hatching from an egg on the moon"],
            ["a cat holding a sign that says hello world"],
            ["an anime illustration of a wiener schnitzel"]
        ], inputs=input_box)

    # Add a button to check the server status
    with gr.Row():
        status_button = gr.Button("Check Server Status")
        status_output = gr.Textbox(label="Server Status")

    def check_status():
        return check_server_status()

    status_button.click(check_status, outputs=status_output)

# Launch the Gradio app
app.launch()