import os
import wget
from stable_diffusion_cpp import StableDiffusion

# Define model paths
vae_url = "https://huggingface.co/pranavajay/flowgram/resolve/main/vae/diffusion_pytorch_model.safetensors"
model_url = "https://huggingface.co/pranavajay/flowgram/resolve/main/flowgram.safetensors"
clip_url = "https://huggingface.co/pranavajay/flowgram/resolve/main/text_encoder/clip_l.safetensors"
t5xxl_url = "https://huggingface.co/pranavajay/flowgram/resolve/main/text_encoder/t5xxl_fp16.safetensors"

# Define local save paths
vae_path = "vae/diffusion_pytorch_model.safetensors"
model_path = "flowgram.safetensors"
clip_path = "text_encoder/clip_l.safetensors"
t5xxl_path = "text_encoder/t5xxl_fp16.safetensors"

# Create directories if they don't exist
os.makedirs("vae", exist_ok=True)
os.makedirs("text_encoder", exist_ok=True)

# Download the models if they do not exist
if not os.path.exists(vae_path):
    wget.download(vae_url, vae_path)
if not os.path.exists(model_path):
    wget.download(model_url, model_path)
if not os.path.exists(clip_path):
    wget.download(clip_url, clip_path)
if not os.path.exists(t5xxl_path):
    wget.download(t5xxl_url, t5xxl_path)

# Initialize the StableDiffusion model
flowgram_diffusion = StableDiffusion(
    diffusion_model_path=model_path,
    clip_l_path=clip_path,
    t5xxl_path=t5xxl_path,
    vae_path=vae_path,
)

# Function to generate an image from text
def generate_image(prompt, num_images=1, guidance_scale=7.5):
    # Generate images
    images = flowgram_diffusion.generate(prompt, num_images=num_images, guidance_scale=guidance_scale)
    
    # Return the generated images
    return images

# Example usage
if __name__ == "__main__":
    prompt = "A beautiful landscape with mountains and a river"
    generated_images = generate_image(prompt)
    
    # Save or display the images
    for i, img in enumerate(generated_images):
        img.save(f"generated_image_{i}.png")  # Save each generated image