File size: 1,608 Bytes
e6773fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
# app.py — Verified Hugging Face Space Code (CPU/GPU Safe)
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import traceback
# ------------------ CONFIG ------------------
MODEL_ID = "runwayml/stable-diffusion-v1-5"
# ------------------ SAFE MODEL LOAD ------------------
try:
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
pipe = StableDiffusionPipeline.from_pretrained(
MODEL_ID,
torch_dtype=dtype,
safety_checker=None
).to(device)
except Exception as e:
print("❌ Model Loading Error:")
traceback.print_exc()
pipe = None
# ------------------ GENERATION FUNCTION ------------------
def generate_image(prompt):
if pipe is None:
return "⚠️ Model not loaded properly. Please check logs."
try:
result = pipe(prompt)
image = result.images[0]
return image
except Exception as e:
print("❌ Generation Error:", e)
return "⚠️ Error generating image."
# ------------------ GRADIO UI ------------------
title = "🪷 Shrividya Text-to-Image AI"
description = "Generate stunning images from your text prompts using Stable Diffusion."
iface = gr.Interface(
fn=generate_image,
inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g. Divine temple on riverbank at sunset"),
outputs=gr.Image(label="Generated Image"),
title=title,
description=description,
allow_flagging="never"
)
if __name__ == "__main__":
iface.launch()
|