Update app.py

app.py

@@ -1,9 +1,7 @@
 import spaces
 import gradio as gr
 import numpy as np
-
 #import tensorrt as trt
-
 import random
 import torch
 from diffusers import StableDiffusion3Pipeline, AutoencoderKL, StableDiffusionXLImg2ImgPipeline, EulerAncestralDiscreteScheduler 
@@ -30,7 +28,6 @@ torch.backends.cudnn.deterministic = False
 #torch.backends.cudnn.benchmark = False
 torch.backends.cuda.preferred_blas_library="cublas"
 #torch.backends.cuda.preferred_linalg_library="cusolver"
-
 torch.set_float32_matmul_precision("highest")
 
 hftoken = os.getenv("HF_AUTH_TOKEN") 
@@ -81,7 +78,6 @@ refiner.scheduler=EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.
 #refiner.scheduler = EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config, beta_schedule="scaled_linear")
 #refiner.scheduler = EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config)
 
-
 tokenizer = AutoTokenizer.from_pretrained(checkpoint, add_prefix_space=False, device_map='balanced')
 tokenizer.tokenizer_legacy=False
 model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map='balanced')
@@ -190,8 +186,8 @@ def infer(
         #sd_image_b = pipe.vae.encode(sd_image_a.to(torch.bfloat16)).latent_dist.sample().mul_(0.18215)
         print("-- using latent file --")
         print('-- generating image --')
-        with torch.no_grad():
-          sd_image = pipe(
+        #with torch.no_grad():
+        sd_image = pipe(
             prompt=enhanced_prompt,  # This conversion is fine
             negative_prompt=negative_prompt,
             guidance_scale=guidance_scale,
@@ -200,11 +196,11 @@ def infer(
             height=height,
             latents=sd_image_a,
             generator=generator
-          ).images[0]
+        ).images[0]
     else:
         print('-- generating image --')
-        with torch.no_grad():
-          sd_image = pipe(
+        #with torch.no_grad():
+        sd_image = pipe(
             prompt=enhanced_prompt,  # This conversion is fine
             prompt_2=enhanced_prompt_2,
             prompt_3=prompt,
@@ -215,7 +211,7 @@ def infer(
             height=height,
          #   latents=None,
             generator=generator,
-          ).images[0]
+        ).images[0]
         print('-- got image --')
     image_path = f"sd35m_{seed}.png"
     sd_image.save(image_path,optimize=False,compress_level=0)
@@ -223,8 +219,8 @@ def infer(
         # Convert the generated image to a tensor
     generated_image_tensor = torch.tensor([np.array(sd_image).transpose(2, 0, 1)]).to('cuda') / 255.0
     # Encode the generated image into latents
-    with torch.no_grad():
-        generated_latents = pipe.vae.encode(generated_image_tensor.to(torch.bfloat16)).latent_dist.sample().mul_(0.18215)
+    #with torch.no_grad():
+    generated_latents = pipe.vae.encode(generated_image_tensor.to(torch.bfloat16)).latent_dist.sample().mul_(0.18215)
     latent_path = f"sd35m_{seed}.pt"
     # Save the latents to a .pt file
     torch.save(generated_latents, latent_path)
@@ -252,9 +248,9 @@ examples = [
 css = """
 #col-container {
     margin: 0 auto;
-    max-width: 768px;
+    max-width: 640px;
 }
-body {
+body{
   background-color: blue;
 }
 """