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

# Initialize the pipelines

text2img_pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
text2img_pipe.to("cuda")
text2img_pipe.load_lora_weights("AbdalrhmanRi/SDXL-Turbo-With-AppleVisionPro", weight_name="pytorch_lora_weights.safetensors")

img2img_pipe = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
img2img_pipe.to("cuda")
img2img_pipe.load_lora_weights("AbdalrhmanRi/SDXL-Turbo-With-AppleVisionPro", weight_name="pytorch_lora_weights.safetensors")


def generate_image(prompt, init_image):
    if init_image is None:

        # Text-to-Image generation
        output_image = text2img_pipe(
            prompt,
            num_inference_steps=40,
            guidance_scale=2.0,
            height=480
        ).images[0]

        yield output_image, None

        output_refiner_image = text2img_pipe(
            prompt=prompt,
            image=output_image,
            guidance_scale=1.0,
            height=480
        ).images[0]

        yield output_image, output_refiner_image


    else:

        # Image-to-Image generation
        init_image = init_image.resize((512, 512))

        output_image = img2img_pipe(
            prompt,
            image=init_image,
            num_inference_steps=40,
            strength=0.5,
            guidance_scale=2.0,
            height=480
        ).images[0]

        yield output_image, None

        output_refiner_image = img2img_pipe(
            prompt=prompt,
            image=output_image,
            strength=0.5,
            guidance_scale=1.0,
            height=480
        ).images[0]

        yield output_image, output_refiner_image

# Define the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt here..."),
        gr.Image(type="pil", label="Initial Image (Optional)", height=360)
    ],
    outputs=[gr.Image(type="pil", label="Generated Image"), gr.Image(type="pil", label="Refined Image")],
    title="Genrate Image Using Generative AI",
    theme=gr.themes.Default(primary_hue="green"),
    description="Text-to-Image or Image-to-Image Generation with SDXL-Turbo."
)

# Launch the interface
interface.launch()