import torch
import gradio as gr
from diffusers import StableDiffusion3Pipeline

def image_generation(prompt):
    # is_cuda = False
    pipeline = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers",
                                                        torch_dtype=torch.float32,
                                                        text_encoder_3=None,
                                                        tokenizer_3=None)
    # pipeline.enable_model_cpu_offload()
    pipeline.to('cpu')

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

    return image

# image_generation("A magicion cat doing spell")

gr.close_all()

interface = gr.Interface(fn=image_generation,
                    inputs=[gr.Textbox(label="Enter your prompt", lines=2)],
                    outputs=[gr.Image(type="pil", label="Generated Image")],
                    title="@GenAILearniverse Project 8: Image creation using Stable Diffusion 3 Model",
                    description="THIS APPLICATION WILL BE USED TO GENERATE AWESOME IMAGES USING SD3 MODEL.")

interface.launch()