import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load the fine-tuned model
def load_model():
    model_name = "MegaTronX/SuicideGirl-FLUX"  # Replace with your model path
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    return model, tokenizer

model, tokenizer = load_model()

# Function to generate an image
def generate_image(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs["input_ids"], max_length=50)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="text",
    title="Image Generation with LoRA Fine-Tuned Model",
    description="Enter a prompt to generate an image using the fine-tuned model."
)

iface.launch()