rtungaturti commited on
Commit
ac7bed8
·
1 Parent(s): a5b9de2

Add app.py

Browse files
Files changed (1) hide show
  1. app.py +99 -0
app.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # %% [markdown]
2
+ # # 🖼️ Tiny Stable Diffusion (CPU Version)
3
+ # **0.9GB Model | No GPU Required**
4
+
5
+ # %% [markdown]
6
+ # ## 1. Install Requirements
7
+ !pip install -q torch diffusers transformers pillow huggingface_hub
8
+ import torch
9
+ from diffusers import StableDiffusionPipeline
10
+ from huggingface_hub import snapshot_download
11
+ from PIL import Image
12
+ import gradio as gr
13
+ import os
14
+
15
+ # Force CPU mode
16
+ torch.backends.quantized.engine = 'qnnpack' # ARM optimization
17
+ device = torch.device("cpu")
18
+
19
+ # %% [markdown]
20
+ # ## 2. Download Model (0.9GB)
21
+ model_path = "./tiny_model"
22
+ os.makedirs(model_path, exist_ok=True)
23
+
24
+ # Download with progress bar
25
+ print("Downloading model... (this may take a few minutes)")
26
+ snapshot_download(
27
+ repo_id="nota-ai/bk-sdm-tiny",
28
+ local_dir=model_path,
29
+ ignore_patterns=["*.bin", "*.fp16*", "*.onnx"],
30
+ local_dir_use_symlinks=False
31
+ )
32
+
33
+ # Verify download
34
+ if not os.listdir(model_path):
35
+ raise ValueError("Model failed to download! Check internet connection")
36
+ else:
37
+ print("✔ Model downloaded successfully")
38
+
39
+ # %% [markdown]
40
+ # ## 3. Load Optimized Pipeline
41
+ print("Loading model...")
42
+ pipe = StableDiffusionPipeline.from_pretrained(
43
+ model_path,
44
+ torch_dtype=torch.float32,
45
+ safety_checker=None,
46
+ requires_safety_checker=False
47
+ ).to(device)
48
+
49
+ # Memory optimizations
50
+ pipe.enable_attention_slicing()
51
+ pipe.unet = torch.compile(pipe.unet) # Compile for faster inference
52
+
53
+ # %% [markdown]
54
+ # ## 4. Generation Function
55
+ def generate_image(prompt, steps=15, seed=42):
56
+ generator = torch.Generator(device).manual_seed(seed)
57
+
58
+ print(f"Generating: {prompt}")
59
+ image = pipe(
60
+ prompt,
61
+ num_inference_steps=steps,
62
+ guidance_scale=7.0,
63
+ generator=generator,
64
+ width=256,
65
+ height=256
66
+ ).images[0]
67
+
68
+ return image
69
+
70
+ # %% [markdown]
71
+ # ## 5. Gradio Interface
72
+ with gr.Blocks(title="Tiny Diffusion (CPU)", css="footer {visibility: hidden}") as demo:
73
+ gr.Markdown("## 🎨 CPU Image Generator (0.9GB Model)")
74
+ with gr.Row():
75
+ prompt = gr.Textbox(label="Prompt",
76
+ value="a cute robot wearing a hat",
77
+ placeholder="Describe your image...")
78
+ with gr.Row():
79
+ steps = gr.Slider(5, 25, value=15, label="Steps")
80
+ seed = gr.Number(42, label="Seed")
81
+ with gr.Row():
82
+ generate_btn = gr.Button("Generate", variant="primary")
83
+ with gr.Row():
84
+ output = gr.Image(label="Output", width=256, height=256)
85
+
86
+ generate_btn.click(
87
+ fn=generate_image,
88
+ inputs=[prompt, steps, seed],
89
+ outputs=output
90
+ )
91
+
92
+ # %% [markdown]
93
+ # ## 6. Launch App
94
+ print("Starting interface...")
95
+ demo.launch(
96
+ server_name="0.0.0.0",
97
+ server_port=7860,
98
+ show_error=True
99
+ )