import torch
from diffusers import AutoPipelineForText2Image
import gradio as gr
import peft

MODEL_NAME = "stabilityai/sdxl-turbo"

pipe = AutoPipelineForText2Image.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, variant="fp16")
pipe.to("cuda")

pipe.load_lora_weights("AbdalrhmanRi/SDXL-Turbo-With-AppleVisionPro", weight_name="pytorch_lora_weights.safetensors")

def generate_image(prompt):
    image = pipe(prompt=prompt, guidance_scale=2.0, num_inference_steps=40, height=480)
    image = image.images[0]

    yield image, None

    refiner_image = pipe(prompt=prompt, image=image, guidance_scale=1.0, height=480)
    refiner_image = refiner_image.images[0]

    yield image, refiner_image


# Set up the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.components.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt here..."),
    outputs=[gr.Image(type="pil", label="Generated Image"), gr.Image(type="pil", label="Refined Image")],
    title="Generate Image Using Generative AI",
    theme=gr.themes.Default(primary_hue="green")
)

# Launch the interface
interface.launch()