hellokawei commited on
Commit
0e69307
·
verified ·
1 Parent(s): 6b22b5b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +30 -0
app.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import StableDiffusionPipeline
4
+ from PIL import Image
5
+ import traceback
6
+ from typing import Optional
7
+
8
+ model_id: str = "runwayml/stable-diffusion-v1-5"
9
+ device: str = "cpu" # force CPU usage for compatibility
10
+
11
+ image_generator_pipe: Optional[StableDiffusionPipeline] = None
12
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
13
+ image_generator_pipe = pipe.to(device)
14
+
15
+
16
+ def generate_image_sd(prompt: str, negative_prompt: str, guidance_scale: float, num_inference_steps: int) -> Image.Image:
17
+ with torch.no_grad():
18
+ output = image_generator_pipe(
19
+ prompt=prompt,
20
+ negative_prompt=negative_prompt,
21
+ guidance_scale=guidance_scale,
22
+ num_inference_steps=num_inference_steps
23
+ )
24
+ image = output.images[0] if output.images else None
25
+
26
+ if not image:
27
+ raise RuntimeError("No image was returned from the generation pipeline.")
28
+
29
+ print(f"Image generated in {end_time - start_time:.2f} seconds.")
30
+ return image