import gradio as gr
import torch
from diffusers import StableDiffusion3Pipeline

def image_generation(prompt):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Load the pipeline (with resume_download if interrupted previously)
    pipeline = StableDiffusion3Pipeline.from_pretrained(
        "stabilityai/stable-diffusion-3-medium-diffusers",
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        text_encoder_3=None,
        tokenizer_3=None,
        resume_download=True
    )

    # ✅ Only use this line if you have GPU + Accelerate
    # pipeline.enable_model_cpu_offload()

    # ✅ Instead, move pipeline to CPU or CUDA manually
    pipeline.to(device)

    image = pipeline(
        prompt=prompt,
        negative_prompt="blurred,ugly,watermark, low resolution, blurry",
        num_inference_steps=50,
        height=1024,
        width=1024,
        guidance_scale=9.0,
    ).images[0]

    return image  # ✅ Return the image for Gradio

# Gradio UI
interface = gr.Interface(
    fn=image_generation,
    inputs=gr.Textbox(lines=2, placeholder="Enter your Prompt..."),
    outputs=gr.Image(type="pil"),
    title="AI Image Generator By Arnav Anand",
    description="This application will be used to generate awesome images using SD3 model"
)

interface.launch()