import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline
from datetime import datetime
import os

class RafayyAI:
    def __init__(self):
        self.model = StableDiffusionXLPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            torch_dtype=torch.float16,
            use_safetensors=True,
            variant="fp16"
        )
        if torch.cuda.is_available():
            self.model = self.model.to("cuda")
        
    def generate_image(self, prompt, negative_prompt=""):
        # Generate unique filename
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"generated_{timestamp}.png"
        
        # Generate image
        image = self.model(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=30,
            guidance_scale=7.5
        ).images[0]
        
        # Save image
        image.save(filename)
        return filename

# Initialize the AI
rafayy = RafayyAI()

# Create Gradio interface
def generate(prompt, negative_prompt=""):
    return rafayy.generate_image(prompt, negative_prompt)

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate..."),
        gr.Textbox(label="Negative Prompt (Optional)", placeholder="What you don't want in the image...")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Rafayy AI Image Generator",
    description="Generate unique images from text descriptions",
    examples=[
        ["A beautiful sunset over mountains", "blur, low quality"],
        ["A futuristic city at night", "dark, blurry"],
        ["A cute cat playing with yarn", "ugly, distorted"]
    ]
)

# Launch the app
if __name__ == "__main__":
    demo.launch()