Update app_with_diffusers.py

app_with_diffusers.py

@@ -39,14 +39,42 @@ pipe.aggregator.load_state_dict(pretrained_state_dict)
 pipe.to(device='cuda', dtype=torch.float16)
 pipe.aggregator.to(device='cuda', dtype=torch.float16)
 
-def infer(prompt, input_image):
+PROMPT = "Photorealistic, highly detailed, hyper detailed photo - realistic maximum detail, 32k, \
+ultra HD, extreme meticulous detailing, skin pore detailing, \
+hyper sharpness, perfect without deformations, \
+taken using a Canon EOS R camera, Cinematic, High Contrast, Color Grading. "
+
+NEG_PROMPT = "blurry, out of focus, unclear, depth of field, over-smooth, \
+sketch, oil painting, cartoon, CG Style, 3D render, unreal engine, \
+dirty, messy, worst quality, low quality, frames, painting, illustration, drawing, art, \
+watermark, signature, jpeg artifacts, deformed, lowres"
+
+def infer(prompt, input_image, steps=30, cfg_scale=7.0, guidance_end=1.0,
+    creative_restoration=False, seed=3407, height=1024, width=1024):
+
+    
     # load a broken image
     low_quality_image = Image.open(input_image).convert("RGB")
+
+    lq = [resize_img(low_quality_image, size=(width, height))]
+    generator = torch.Generator(device=device).manual_seed(seed)
+    timesteps = [
+        i * (1000//steps) + pipe.scheduler.config.steps_offset for i in range(0, steps)
+    ]
+    timesteps = timesteps[::-1]
+
+    prompt = PROMPT if len(prompt)==0 else prompt
+    neg_prompt = NEG_PROMPT
     
     # InstantIR restoration
     image = pipe(
-        prompt=prompt,
-        image=low_quality_image,
+        prompt=[prompt]*len(lq),
+        image=lq,
+        num_inference_steps=steps,
+        generator=generator,
+        timesteps=timesteps,
+        negative_prompt=[neg_prompt]*len(lq),
+        guidance_scale=cfg_scale,
         previewer_scheduler=lcm_scheduler,
     ).images[0]
 
@@ -54,12 +82,16 @@ def infer(prompt, input_image):
 
 import gradio as gr
 
+
+
 with gr.Blocks() as demo:
     with gr.Column():
         with gr.Row():
             with gr.Column():
                 lq_img = gr.Image(label="Low-quality image", type="filepath")
-                prompt = gr.Textbox(label="Prompt", value="")
+                with gr.Group():
+                    prompt = gr.Textbox(label="Prompt", value="")
+                    
                 submit_btn = gr.Button("InstantIR magic!")
             output_img = gr.Image(label="InstantIR restored")
     submit_btn.click(