Update app.py

app.py

@@ -17,6 +17,8 @@ import torch
 import cv2
 import gc
 
+from image_gen_aux import UpscaleWithModel
+
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
@@ -40,6 +42,8 @@ import shutil
 
 MAX_SEED = np.iinfo(np.int32).max
 
+upscaler = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))
+
 #import diffusers
 from diffusers import StableDiffusionXLImg2ImgPipeline, AutoencoderKL
 print("Loading SDXL Image-to-Image pipeline...")
@@ -117,16 +121,20 @@ def get_duration(*args, **kwargs):
     return 90
 
 
-@spaces.GPU()
-def enhance_frame(image_to_enhance: Image.Image):
+@spaces.GPU(duration=45)
+def enhance_frame(prompt, image_to_enhance: Image.Image):
     try:
         print("Moving enhancer pipeline to GPU...")
         seed = random.randint(0, MAX_SEED)
         generator = torch.Generator(device='cpu').manual_seed(seed)
         enhancer_pipeline.to("cuda",torch.bfloat16)
-        refine_prompt = "cinematic, high detail, sharp focus, 8k, professional photography"
-        enhanced_image = enhancer_pipeline(prompt=refine_prompt, image=image_to_enhance, strength=0.125, generator=generator, num_inference_steps=100).images[0]
+        refine_prompt = prompt +" high detail, sharp focus, 8k, professional"
+        enhanced_image = enhancer_pipeline(prompt=refine_prompt, image=image_to_enhance, strength=0.1, generator=generator, num_inference_steps=220).images[0]
         print("Frame enhancement successful.")
+        print("Doing super-resolution.")
+        with torch.no_grad():
+            upscale = upscaler(enhanced_image, tiling=True, tile_width=1024, tile_height=1024)
+            enhanced_image = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
     except Exception as e:
         print(f"Error during frame enhancement: {e}")
         gr.Warning("Frame enhancement failed. Using original frame.")
@@ -139,7 +147,7 @@ def enhance_frame(image_to_enhance: Image.Image):
     return enhanced_image
 
 
-def use_last_frame_as_input(video_filepath, do_enhance):
+def use_last_frame_as_input(prompt, video_filepath, do_enhance):
     if not video_filepath or not os.path.exists(video_filepath):
         gr.Warning("No video clip available.")
         return None, gr.update()
@@ -155,7 +163,7 @@ def use_last_frame_as_input(video_filepath, do_enhance):
         print("Displaying original last frame...")
         yield pil_image, gr.update()
         if do_enhance:
-            enhanced_image = enhance_frame(pil_image)
+            enhanced_image = enhance_frame(prompt, pil_image)
             # 2. Yield the enhanced frame and switch the tab
             print("Displaying enhanced frame and switching tab...")
             yield enhanced_image, gr.update(selected="i2v_tab")
@@ -312,7 +320,7 @@ with gr.Blocks(css=css) as demo:
     t2v_button.click(hide_btn, outputs=[use_last_frame_button], queue=False).then(fn=generate, inputs=t2v_inputs, outputs=gen_outputs, api_name="text_to_video")
     i2v_button.click(hide_btn, outputs=[use_last_frame_button], queue=False).then(fn=generate, inputs=i2v_inputs, outputs=gen_outputs, api_name="image_to_video")
     v2v_button.click(hide_btn, outputs=[use_last_frame_button], queue=False).then(fn=generate, inputs=v2v_inputs, outputs=gen_outputs, api_name="video_to_video")
-    use_last_frame_button.click(fn=use_last_frame_as_input, inputs=[output_video,enhance_checkbox], outputs=[image_i2v, tabs])
+    use_last_frame_button.click(fn=use_last_frame_as_input, inputs=[i2v_prompt,output_video,enhance_checkbox], outputs=[image_i2v, tabs])
     stitch_button.click(fn=stitch_videos, inputs=[clips_state], outputs=[final_video_output])
     clear_button.click(fn=clear_clips, outputs=[clips_state, clip_counter_display, output_video, final_video_output])
 if __name__ == "__main__":