import torch
from diffusers import StableDiffusion3Pipeline
from huggingface_hub import login
import os
import gradio as gr

# Check for GPU availability and set device accordingly
if torch.cuda.is_available():
    device = "cuda"
    print("GPU is available")
else:
    device = "cpu"
    print("GPU is not available, using CPU")

# Retrieve the token from the environment variable
token = os.getenv("HF_TOKEN")  # Hugging Face token from the secret
if token:
    login(token=token)  # Log in with the retrieved token
else:
    raise ValueError("Hugging Face token not found. Please set it as a repository secret in the Space settings.")

# Load the Stable Diffusion 3.5 model with lower precision (float16) if GPU is available
model_id = "stabilityai/stable-diffusion-3.5-large"
if device == "cuda":
    pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)  # Use float16 precision
else:
    pipe = StableDiffusion3Pipeline.from_pretrained(model_id)  # Default precision for CPU

pipe.to(device)  # Ensuring the model is on the correct device (GPU or CPU)

# Define the path to the LoRA model
lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally

# Custom method to load and apply LoRA weights to the Stable Diffusion pipeline
def load_lora_model(pipe, lora_model_path):
    # Load the LoRA weights
    lora_weights = torch.load(lora_model_path, map_location=device)  # Load LoRA model to the correct device
    
    # Print available attributes of the model to check access to `unet` (optional)
    print(dir(pipe))  # This will list all attributes and methods of the `pipe` object

    # Apply weights to the UNet submodule
    try:
        for name, param in pipe.unet.named_parameters():  # Accessing unet parameters
            if name in lora_weights:
                param.data += lora_weights[name]
    except AttributeError:
        print("The model doesn't have 'unet' attributes. Please check the model structure.")
        # Add alternative handling or exit

    return pipe

# Load and apply the LoRA model weights
pipe = load_lora_model(pipe, lora_model_path)

# Function to generate an image from a text prompt
def generate_image(prompt, seed=None):
    generator = torch.manual_seed(seed) if seed is not None else None
    # Reduce image size for less memory usage
    image = pipe(prompt, height=512, width=512, generator=generator).images[0]  # Changed image size
    return image

# Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Enter your prompt"),  # For the prompt
        gr.Number(label="Enter a seed (optional)", value=None),  # For the seed
    ],
    outputs="image"
)
iface.launch()