Update app.py

app.py

@@ -6,16 +6,15 @@ import torch
 from diffusers import DiffusionPipeline
 from transformers import MarianMTModel, MarianTokenizer
 
-dtype = torch.float32  # Zmiana na float32 dla lepszej kompatybilności z CPU
-device = "cpu"  # Wymuszenie użycia CPU
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load the translation model and tokenizer
 model_name = 'Helsinki-NLP/opus-mt-pl-en'
 tokenizer = MarianTokenizer.from_pretrained(model_name)
 model = MarianMTModel.from_pretrained(model_name).to(device)
 
-# Load the new diffusion model
-pipe = DiffusionPipeline.from_pretrained("XLabs-AI/flux-RealismLora", torch_dtype=dtype).to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
@@ -294,4 +293,4 @@ with gr.Blocks(theme=gr.themes.Monochrome(), css=css) as demo:
         api_name=False,
     )
 
-demo.launch()
+demo.launch()