app.py

import torch
from diffusers import DiffusionPipeline
import gradio as gr

# Load model with float16 for GPU or float32 for CPU
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", torch_dtype=dtype
)
pipe.load_lora_weights("EliKet/train_text_to_img")
pipe.to(device)

def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(placeholder="Enter your image prompt here..."),
    outputs="image",
    title="Text-to-Image Generator (Lynx)",
    description="Type a prompt like 'a lynx in the snowy forest, ultra-detailed'."
)

demo.launch()