Alibrown commited on
Commit
26f9444
·
verified ·
1 Parent(s): 50374e6

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +94 -0
app.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import random
3
+ import numpy as np
4
+ import torch
5
+ import gradio as gr
6
+ from diffusers import DiffusionPipeline # Geändert
7
+ import paramiko
8
+ from huggingface_hub import login
9
+
10
+ HF_TOKEN = os.getenv('HF_TOKEN', '').strip()
11
+ if not HF_TOKEN:
12
+ raise ValueError("HUGGING_TOKEN is not set.")
13
+
14
+ login(token=HF_TOKEN)
15
+
16
+ # Konfiguration
17
+ STORAGE_DOMAIN = os.getenv('STORAGE_DOMAIN', '').strip() # SFTP Server Domain
18
+ STORAGE_USER = os.getenv('STORAGE_USER', '').strip() # SFTP User
19
+ STORAGE_PSWD = os.getenv('STORAGE_PSWD', '').strip() # SFTP Passwort
20
+ STORAGE_PORT = int(os.getenv('STORAGE_PORT', '22').strip()) # SFTP Port
21
+ STORAGE_SECRET = os.getenv('STORAGE_SECRET', '').strip() # Secret Token
22
+
23
+ # Modell laden mit DiffusionPipeline
24
+ device = "cuda" if torch.cuda.is_available() else "cpu"
25
+ repo = "stabilityai/stable-diffusion-3-medium-diffusers"
26
+
27
+ try:
28
+ pipe = DiffusionPipeline.from_pretrained( # Geändert
29
+ repo,
30
+ torch_dtype=torch.float16,
31
+ use_safetensors=True
32
+ ).to(device)
33
+ except Exception as e:
34
+ raise RuntimeError(f"Model load error: {e}")
35
+
36
+ # Maximalwerte
37
+ MAX_SEED = np.iinfo(np.int32).max
38
+ MAX_IMAGE_SIZE = 1344
39
+
40
+ # SFTP-Funktion
41
+ def upload_to_sftp(local_file, remote_path):
42
+ try:
43
+ transport = paramiko.Transport((STORAGE_DOMAIN, STORAGE_PORT))
44
+ transport.connect(username=STORAGE_USER, password=STORAGE_PSWD)
45
+ sftp = paramiko.SFTPClient.from_transport(transport)
46
+ sftp.put(local_file, remote_path)
47
+ sftp.close()
48
+ transport.close()
49
+ print(f"File {local_file} successfully uploaded to {remote_path}")
50
+ return True
51
+ except Exception as e:
52
+ print(f"Error during SFTP upload: {e}")
53
+ return False
54
+
55
+ # Inferenz-Funktion
56
+ def infer(prompt, width, height, guidance_scale, num_inference_steps, seed, randomize_seed):
57
+ if randomize_seed:
58
+ seed = random.randint(0, MAX_SEED)
59
+
60
+ generator = torch.manual_seed(seed)
61
+ image = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, generator=generator).images[0]
62
+
63
+ # Speichere Bild lokal
64
+ local_file = f"/tmp/generated_image_{seed}.png"
65
+ image.save(local_file)
66
+
67
+ # Hochladen zu SFTP
68
+ remote_path = f"/uploads/generated_image_{seed}.png"
69
+ if upload_to_sftp(local_file, remote_path):
70
+ os.remove(local_file)
71
+ return f"Image uploaded to {remote_path}", seed
72
+ else:
73
+ return "Failed to upload image", seed
74
+
75
+ # Gradio-App
76
+ with gr.Blocks() as demo:
77
+ gr.Markdown("### Stable Diffusion 3 - Test App")
78
+ prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
79
+ width = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=512, label="Width")
80
+ height = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=512, label="Height")
81
+ guidance_scale = gr.Slider(0.0, 10.0, step=0.1, value=7.5, label="Guidance Scale")
82
+ num_inference_steps = gr.Slider(1, 50, step=1, value=25, label="Inference Steps")
83
+ seed = gr.Number(value=42, label="Seed")
84
+ randomize_seed = gr.Checkbox(value=False, label="Randomize Seed")
85
+ generate_button = gr.Button("Generate Image")
86
+ output = gr.Text(label="Output")
87
+
88
+ generate_button.click(
89
+ infer,
90
+ inputs=[prompt, width, height, guidance_scale, num_inference_steps, seed, randomize_seed],
91
+ outputs=[output, seed]
92
+ )
93
+
94
+ demo.launch()