Update app.py

app.py

@@ -8,7 +8,6 @@ from PIL import Image, ImageDraw, ImageFont
 import gradio as gr
 from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
 from diffusers.utils.torch_utils import randn_tensor
-from tqdm import tqdm
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -168,7 +167,7 @@ class StableDiffusionPipeline:
         shape = (1, self.vae.config.latent_channels, mask_height // vae_scale_factor, mask_width // vae_scale_factor)
         latents = randn_tensor(shape, generator=torch.manual_seed(20), device=self.device) * self.scheduler.init_noise_sigma
 
-        for t in tqdm(timesteps):
+        for t in timesteps:
             latent_model_input = latents
             latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
             sample = torch.cat([latent_model_input, masked_image_latents, glyph_latents, mask], dim=1)
@@ -181,8 +180,9 @@ class StableDiffusionPipeline:
         return image, image_vae
 
 # Load models (adjust the paths to your model directories)
-vae = AutoencoderKL.from_pretrained("Yesianrohn/TextSSR", subfolder="vae")
-unet = UNet2DConditionModel.from_pretrained("Yesianrohn/TextSSR", subfolder="unet")
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+vae = AutoencoderKL.from_pretrained("Yesianrohn/TextSSR", subfolder="vae", torch_dtype=dtype)
+unet = UNet2DConditionModel.from_pretrained("Yesianrohn/TextSSR", subfolder="unet", torch_dtype=dtype)
 noise_scheduler = DDPMScheduler.from_pretrained("Yesianrohn/TextSSR", subfolder="scheduler")
 
 # Create pipeline