Update app.py

app.py

@@ -16,35 +16,25 @@ if not HF_TOKEN:
 login(token=HF_TOKEN)
 
 # Konfiguration
-STORAGE_DOMAIN = os.getenv('STORAGE_DOMAIN', '').strip()  # SFTP Server Domain
-STORAGE_USER = os.getenv('STORAGE_USER', '').strip()  # SFTP User
-STORAGE_PSWD = os.getenv('STORAGE_PSWD', '').strip()  # SFTP Passwort
-STORAGE_PORT = int(os.getenv('STORAGE_PORT', '22').strip())  # SFTP Port
-STORAGE_SECRET = os.getenv('STORAGE_SECRET', '').strip()  # Secret Token
+STORAGE_DOMAIN = os.getenv('STORAGE_DOMAIN', '').strip() # SFTP Server Domain
+STORAGE_USER = os.getenv('STORAGE_USER', '').strip() # SFTP User
+STORAGE_PSWD = os.getenv('STORAGE_PSWD', '').strip() # SFTP Passwort
+STORAGE_PORT = int(os.getenv('STORAGE_PORT', '22').strip()) # SFTP Port
+STORAGE_SECRET = os.getenv('STORAGE_SECRET', '').strip() # Secret Token
 
-# Modell-Konfiguration und Device-Setup
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
-
-# Stelle fest, ob auf CPU oder GPU-System
-is_gpu_available = torch.cuda.is_available()
+# Modell-Optionen - können angepasst werden
+MODEL_REPO = os.getenv('MODEL_REPO', 'stabilityai/stable-diffusion-2-1-base')  # Standard-Modell
+TORCH_DTYPE = os.getenv('TORCH_DTYPE', 'float16')  # Standard-Präzision
 
-# Modell laden - passend zur Hardware
-repo = "stabilityai/stable-diffusion-3-medium-diffusers"
-
-# Die Standard-Präzision basiert auf verfügbarer Hardware
-DEFAULT_PRECISION = "float16" if is_gpu_available else "float32"
-print(f"Default precision: {DEFAULT_PRECISION}")
+# Modell laden
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if TORCH_DTYPE == 'float16' else torch.float32
 
-# Modell beim Start laden
 try:
-    # Wähle Präzision basierend auf Hardware
-    if DEFAULT_PRECISION == "float16":
-        pipe = StableDiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16).to(device)
-    else:  # float32 für CPU
-        pipe = StableDiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float32).to(device)
-    
-    print("Model loaded successfully")
+    pipe = StableDiffusionPipeline.from_pretrained(
+        MODEL_REPO, 
+        torch_dtype=torch_dtype
+    ).to(device)
 except Exception as e:
     raise RuntimeError(f"Failed to load the model. Ensure the token has access to the repo. Error: {e}")
 
@@ -94,28 +84,25 @@ def infer(prompt, width, height, guidance_scale, num_inference_steps, seed, rand
     else:
         return "Failed to upload image", seed
 
+# App-Titel mit Modell- und Präzisionsinformationen
+APP_TITLE = f"### Stable Diffusion - {os.path.basename(MODEL_REPO)} ({TORCH_DTYPE} auf {device})"
+
 # Gradio-App
 with gr.Blocks() as demo:
-    gr.Markdown(f"### Stable Diffusion 3 - Test App (Running on {device.upper()} with {DEFAULT_PRECISION})")
-    
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
-            width = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=512, label="Width")
-            height = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=512, label="Height")
-            guidance_scale = gr.Slider(0.0, 10.0, step=0.1, value=7.5, label="Guidance Scale")
-            num_inference_steps = gr.Slider(1, 50, step=1, value=25, label="Inference Steps")
-            seed = gr.Number(value=42, label="Seed")
-            randomize_seed = gr.Checkbox(value=False, label="Randomize Seed")
-            generate_button = gr.Button("Generate Image")
-            output = gr.Text(label="Output")
+    gr.Markdown(APP_TITLE)
+    prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
+    width = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=512, label="Width")
+    height = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=512, label="Height")
+    guidance_scale = gr.Slider(0.0, 10.0, step=0.1, value=7.5, label="Guidance Scale")
+    num_inference_steps = gr.Slider(1, 50, step=1, value=25, label="Inference Steps")
+    seed = gr.Number(value=42, label="Seed")
+    randomize_seed = gr.Checkbox(value=False, label="Randomize Seed")
+    generate_button = gr.Button("Generate Image")
+    output = gr.Text(label="Output")
     
     generate_button.click(
         infer,
-        inputs=[
-            prompt, width, height, guidance_scale, 
-            num_inference_steps, seed, randomize_seed
-        ],
+        inputs=[prompt, width, height, guidance_scale, num_inference_steps, seed, randomize_seed],
         outputs=[output, seed]
     )