Update app.py

app.py

@@ -53,20 +53,20 @@ torch_dtype = torch.bfloat16
 
 checkpoint = "microsoft/Phi-3.5-mini-instruct"
 #vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-#vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", torch_dtype=torch.bfloat16)
+vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", torch_dtype=torch.bfloat16, device_map='balanced')
 
-pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16", torch_dtype=torch.bfloat16)
+pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16", torch_dtype=torch.bfloat16, device_map='balanced')
 #pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", token=hftoken, torch_dtype=torch.float32, device_map='balanced')
 
 # pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
 
 #pipe.scheduler.config.requires_aesthetics_score = False
 #pipe.enable_model_cpu_offload()
-pipe.to(device)
+#pipe.to(device)
 #pipe = torch.compile(pipe)
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear", algorithm_type="sde-dpmsolver++")
+# pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear")
 
-#refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("ford442/stable-diffusion-xl-refiner-1.0-bf16", vae=vae, torch_dtype=torch.bfloat16, use_safetensors=True, requires_aesthetics_score=True, device_map='balanced')
+refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("ford442/stable-diffusion-xl-refiner-1.0-bf16", vae=vae, torch_dtype=torch.bfloat16, use_safetensors=True, requires_aesthetics_score=True, device_map='balanced')
 #refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float32, requires_aesthetics_score=True, device_map='balanced')
 
 #refiner.enable_model_cpu_offload()
@@ -74,7 +74,7 @@ pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.conf
 #refiner.scheduler.config.requires_aesthetics_score=False
 #refiner.to(device)
 #refiner = torch.compile(refiner)
-#refiner.scheduler = EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config, beta_schedule="scaled_linear")
+refiner.scheduler = EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config, beta_schedule="scaled_linear")
 
 tokenizer = AutoTokenizer.from_pretrained(checkpoint, add_prefix_space=False, device_map='balanced')
 tokenizer.tokenizer_legacy=False
@@ -90,7 +90,7 @@ def filter_text(text):
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 4096
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=80)
 def infer(
     prompt,
     negative_prompt,
@@ -139,21 +139,31 @@ def infer(
     print('-- filtered prompt --')
     print(enhanced_prompt)
     print('-- generating image --')
-    #with torch.no_grad():
     sd_image = pipe(
-            prompt=enhanced_prompt,  # This conversion is fine
-            negative_prompt=negative_prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator
+        prompt=enhanced_prompt,  # This conversion is fine
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator
     ).images[0]
     print('-- got image --')
     image_path = f"sd35m_{seed}.png"
     sd_image.save(image_path,optimize=False,compress_level=0)
     upload_to_ftp(image_path)
-    return sd_image, seed, image_path, enhanced_prompt
+    refine = refiner(
+            prompt=f"{prompt}, high quality masterpiece, complex details",
+            negative_prompt = negative_prompt,
+            guidance_scale=7.5,
+            num_inference_steps=num_inference_steps,
+            image=sd_image,
+            generator=generator,
+    ).images[0]  
+    refine_path = f"refine_{seed}.png"
+    refine.save(refine_path,optimize=False,compress_level=0)
+    upload_to_ftp(refine_path)
+    return refine, seed, refine_path, enhanced_prompt
 
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
@@ -168,6 +178,37 @@ css = """
 }
 """
 
+def repeat_infer(
+    prompt,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    width,
+    height,
+    guidance_scale,
+    num_inference_steps,
+    num_iterations,  # New input for number of iterations
+):
+    i = 0
+    while i < num_iterations:
+        time.sleep(700)  # Wait for 10 minutes (600 seconds)
+        result, seed, image_path, enhanced_prompt = infer(
+            prompt,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+        )
+        
+        # Optionally, you can add logic here to process the results of each iteration
+        # For example, you could display the image, save it with a different name, etc.
+        i += 1
+    return result, seed, image_path, enhanced_prompt
+
+
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(" # Text-to-Text-to-Image StableDiffusion 3.5 Medium (with refine)")
@@ -191,6 +232,9 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter a negative prompt",
                 visible=False,
             )
+            num_iterations = gr.Number(
+                value=1000, 
+                label="Number of Iterations")
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,