ma4389 commited on
Commit
4106559
·
verified ·
1 Parent(s): e74e40f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -27
app.py CHANGED
@@ -1,27 +1,27 @@
1
- # gradio_app.py
2
-
3
- import torch
4
- from diffusers import DiffusionPipeline
5
- import gradio as gr
6
-
7
- # Load model
8
- pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
9
- pipe.to("cuda")
10
- pipe.load_lora_weights("EliKet/train_text_to_img")
11
-
12
- # Generation function
13
- def generate_image(prompt):
14
- image = pipe(prompt).images[0]
15
- return image
16
-
17
- # Gradio Interface
18
- demo = gr.Interface(
19
- fn=generate_image,
20
- inputs=gr.Textbox(lines=2, placeholder="Enter your image prompt here..."),
21
- outputs="image",
22
- title="🐾 Lynx Text-to-Image Generator",
23
- description="Type a prompt (e.g., 'A majestic lynx in a snowy forest') and get an AI-generated image using Stable Diffusion + LoRA."
24
- )
25
-
26
- # Launch app
27
- demo.launch()
 
1
+ import torch
2
+ from diffusers import DiffusionPipeline
3
+ import gradio as gr
4
+
5
+ # Load model with float16 for GPU or float32 for CPU
6
+ dtype = torch.float16 if torch.cuda.is_available() else torch.float32
7
+ device = "cuda" if torch.cuda.is_available() else "cpu"
8
+
9
+ pipe = DiffusionPipeline.from_pretrained(
10
+ "CompVis/stable-diffusion-v1-4", torch_dtype=dtype
11
+ )
12
+ pipe.load_lora_weights("EliKet/train_text_to_img")
13
+ pipe.to(device)
14
+
15
+ def generate_image(prompt):
16
+ image = pipe(prompt).images[0]
17
+ return image
18
+
19
+ demo = gr.Interface(
20
+ fn=generate_image,
21
+ inputs=gr.Textbox(placeholder="Enter your image prompt here..."),
22
+ outputs="image",
23
+ title="Text-to-Image Generator (Lynx)",
24
+ description="Type a prompt like 'a lynx in the snowy forest, ultra-detailed'."
25
+ )
26
+
27
+ demo.launch()